Tianshou

The reinforcement learning process in Tianshou involves the interaction between an agent and an environment. The key components and their Tianshou class correspondences are:

• Environment: The system the agent interacts with. In Tianshou, this is represented by a class inheriting from gymnasium.Env. Environments are often vectorized for parallel interaction.
• Policy: The strategy for action selection. Encapsulated in the tianshou.algorithm.algorithm_base.Policy class.
• Replay Buffer: A data structure storing experiences (state transitions, actions, rewards). Implemented via tianshou.data.buffer.buffer_base.ReplayBuffer. Experiences are added using a tianshou.data.collector.Collector and sampled as tianshou.data.batch.Batch objects.
• Learning Algorithm: Defines how the policy is updated using buffer data. Abstracted by the tianshou.algorithm.algorithm_base.Algorithm class.

The procedural API requires manual orchestration of the RL pipeline.

Workflow:

1. Setup Logging: Use ts.utils.TensorboardLogger.
2. Create Environments: Use ts.env.DummyVectorEnv or similar to wrap Gymnasium environments.
3. Build Network: Extract SpaceInfo from the environment to define state_shape and action_shape, then instantiate a network (e.g., ts.utils.net.common.Net).
4. Create Policy & Algorithm: Instantiate a policy (e.g., DiscreteQLearningPolicy) and pass it to an algorithm (e.g., ts.algorithm.DQN).
5. Setup Collectors: Create ts.data.Collector instances for both training (with a VectorReplayBuffer) and testing.
6. Train: Call algorithm.run_training() passing an OffPolicyTrainerParams object containing collectors, hyperparameters, and a stop_fn.
7. Evaluate: Manually run a collector on a single environment to watch the agent.

import gymnasium as gym
import tianshou as ts
from tianshou.algorithm.modelfree.dqn import DiscreteQLearningPolicy
from tianshou.algorithm.optim import AdamOptimizerFactory
from tianshou.data import CollectStats
from tianshou.trainer import OffPolicyTrainerParams
from tianshou.utils.net.common import Net
from tianshou.utils.space_info import SpaceInfo
from torch.utils.tensorboard import SummaryWriter

# Define hyperparameters
task = "CartPole-v1"
lr, epoch, batch_size = 1e-3, 10, 64
num_training_envs, num_test_envs = 10, 100
gamma, n_step, target_freq = 0.9, 3, 320
buffer_size = 20000
eps_train, eps_test = 0.1, 0.05
epoch_num_steps, collection_step_num_env_steps = 10000, 10

# Set up logging
logger = ts.utils.TensorboardLogger(SummaryWriter("log/dqn"))

# Create environments
training_envs = ts.env.DummyVectorEnv([lambda: gym.make(task) for _ in range(num_training_envs)])
test_envs = ts.env.DummyVectorEnv([lambda: gym.make(task) for _ in range(num_test_envs)])

# Build the network
env = gym.make(task, render_mode="human")
space_info = SpaceInfo.from_env(env)
state_shape = space_info.observation_info.obs_shape
action_shape = space_info.action_info.action_shape
net = Net(state_shape=state_shape, action_shape=action_shape, hidden_sizes=[128, 128, 128])

# Create policy and algorithm
policy = DiscreteQLearningPolicy(
    model=net,
    action_space=env.action_space,
    eps_training=eps_train,
    eps_inference=eps_test,
)
algorithm = ts.algorithm.DQN(
    policy=policy,
    optim=AdamOptimizerFactory(lr=lr),
    gamma=gamma,
    n_step_return_horizon=n_step,
    target_update_freq=target_freq,
)

# Set up collectors
training_collector = ts.data.Collector[CollectStats](
    algorithm,
    training_envs,
    ts.data.VectorReplayBuffer(buffer_size, num_training_envs),
    exploration_noise=True,
)
test_collector = ts.data.Collector[CollectStats](
    algorithm,
    test_envs,
    exploration_noise=True,
)


# Define stop condition
def stop_fn(mean_rewards: float) -> bool:
    if env.spec and env.spec.reward_threshold:
        return mean_rewards >= env.spec.reward_threshold
    return False


# Train the algorithm
result = algorithm.run_training(
    OffPolicyTrainerParams(
        training_collector=training_collector,
        test_collector=test_collector,
        max_epochs=epoch,
        epoch_num_steps=epoch_num_steps,
        collection_step_num_env_steps=collection_step_num_env_steps,
        test_step_num_episodes=num_test_envs,
        batch_size=batch_size,
        update_step_num_gradient_steps_per_sample=1 / collection_step_num_env_steps,
        stop_fn=stop_fn,
        logger=logger,
        test_in_training=True,
    )
)
print(f"Finished training in {result.timing.total_time} seconds")

# Watch the trained agent
collector = ts.data.Collector[CollectStats](algorithm, env, exploration_noise=True)
collector.collect(n_episode=100, render=1 / 35)

To support a new algorithm in the high-level API, you must implement three core components:

1. Parameter Class

Define a dataclass in tianshou/highlevel/params/algorithm_params.py inheriting from Params.

• Choose a base class based on architecture (e.g., actor-critic) and paradigm (on-policy/off-policy).
• Override _get_param_transformers() to define how high-level parameters transform into low-level policy parameters.

2. Algorithm Factory

Implement a subclass of AlgorithmFactory in tianshou/highlevel/algorithm.py.

• Inherit from existing base factories (e.g., ActorCriticOnPolicyAlgorithmFactory) where possible.
• Implement _get_algorithm_class() to return the target algorithm class.
• Override _create_algorithm() or _create_kwargs() if custom instantiation logic is required.

3. Experiment Builder

Add a builder class in tianshou/highlevel/experiment.py inheriting from OnPolicyExperimentBuilder or OffPolicyExperimentBuilder.

• Use mixins for common patterns (e.g., actor/critic configuration).
• Implement _create_algorithm_factory() to instantiate the algorithm factory.
• Optionally add with_* methods for fluent configuration.

Final Step: Export the new classes in tianshou/highlevel/__init__.py.

The high-level API uses ExperimentBuilder classes to construct a complete training pipeline from configuration objects.

Workflow:

1. Define environment configuration using EnvFactoryRegistered.
2. Define experiment settings using ExperimentConfig (persistence, logging, watching).
3. Define training configuration using specific config classes like OffPolicyTrainingConfig.
4. Chain builder methods to add algorithm-specific parameters (e.g., .with_dqn_params()), model architectures (.with_model_factory_default()), and stop conditions (.with_epoch_stop_callback()).
5. Call .build() to create the experiment and .run() to execute it.

from tianshou.highlevel.config import OffPolicyTrainingConfig
from tianshou.highlevel.env import EnvFactoryRegistered, VectorEnvType
from tianshou.highlevel.experiment import DQNExperimentBuilder, ExperimentConfig
from tianshou.highlevel.params.algorithm_params import DQNParams
from tianshou.highlevel.trainer import EpochStopCallbackRewardThreshold

# Build the experiment through configuration
experiment = (
    DQNExperimentBuilder(
        # Environment configuration
        EnvFactoryRegistered(
            task="CartPole-v1",
            venv_type=VectorEnvType.DUMMY,
            training_seed=0,
            test_seed=10,
        ),
        # Experiment settings
        ExperimentConfig(
            persistence_enabled=False,
            watch=True,
            watch_render=1 / 35,
            watch_num_episodes=100,
        ),
        # Training configuration
        OffPolicyTrainingConfig(
            max_epochs=10,
            epoch_num_steps=10000,
            batch_size=64,
            num_training_envs=10,
            num_test_envs=100,
            buffer_size=20000,
            collection_step_num_env_steps=10,
            update_step_num_gradient_steps_per_sample=1 / 10,
        ),
    )
    # Algorithm-specific parameters
    .with_dqn_params(
        DQNParams(
            lr=1e-3,
            gamma=0.9,
            n_step_return_horizon=3,
            target_update_freq=320,
            eps_training=0.3,
            eps_inference=0.0,
        ),
    )
    # Network architecture
    .with_model_factory_default(hidden_sizes=(64, 64))
    # Stop condition
    .with_epoch_stop_callback(EpochStopCallbackRewardThreshold(195))
    .build()
)

# Run the experiment
experiment.run()

The Procedural API provides granular control for advanced users and algorithm developers. The workflow involves:

1. Environment Setup: Create training and test environments using ts.env.DummyVectorEnv or SubprocVectorEnv.
2. Network & Policy: Define a neural network (e.g., using tianshou.utils.net.common.Net) and wrap it in a policy (e.g., DiscreteQLearningPolicy).
3. Algorithm: Initialize the algorithm (e.g., DQN) with the policy and an optimizer factory.
4. Collectors: Set up ts.data.Collector instances for training (using a VectorReplayBuffer) and testing.
5. Training: Execute training using algorithm.run_training() with OffPolicyTrainerParams.
6. Evaluation: Use a Collector to watch the agent's performance.

import gymnasium as gym
import tianshou as ts
from tianshou.algorithm.modelfree.dqn import DiscreteQLearningPolicy
from tianshou.algorithm.optim import AdamOptimizerFactory
from tianshou.data import CollectStats
from tianshou.trainer import OffPolicyTrainerParams
from tianshou.utils.net.common import Net
from tianshou.utils.space_info import SpaceInfo
from torch.utils.tensorboard import SummaryWriter

# ... (Hyper-parameter and environment setup omitted for brevity in this summary)

# Create the network
env = gym.make(task, render_mode="human")
space_info = SpaceInfo.from_env(env)
state_shape = space_info.observation_info.obs_shape
action_shape = space_info.action_info.action_shape
net = Net(state_shape=state_shape, action_shape=action_shape, hidden_sizes=[128, 128, 128])

# Create the policy
policy = DiscreteQLearningPolicy(
    model=net,
    action_space=env.action_space,
    eps_training=eps_train,
    eps_inference=eps_test
)

# Create the algorithm
algorithm = DQN(
    policy=policy,
    optim=AdamOptimizerFactory(lr=lr),
    gamma=gamma,
    n_step_return_horizon=n_step,
    target_update_freq=target_freq
)

# Set up collectors
training_collector = ts.data.Collector[CollectStats](
  algorithm,
  training_envs,
  ts.data.VectorReplayBuffer(buffer_size, num_training_envs),
  exploration_noise=True,
)
test_collector = ts.data.Collector[CollectStats](
  algorithm,
  test_envs,
  exploration_noise=True,
) 

# Train
result = algorithm.run_training(
  OffPolicyTrainerParams(
    training_collector=training_collector,
    test_collector=test_collector,
    max_epochs=epoch,
    epoch_num_steps=epoch_num_steps,
    collection_step_num_env_steps=collection_step_num_env_steps,
    test_step_num_episodes=num_test_envs,
    batch_size=batch_size,
    update_step_num_gradient_steps_per_sample=1 / collection_step_num_env_steps,
    stop_fn=lambda mean_rewards: mean_rewards >= env.spec.reward_threshold,
    logger=logger,
    test_in_training=True,
  )
)