common#


miniblock(input_size: int, output_size: int = 0, norm_layer: type[~torch.nn.modules.module.Module] | None = None, norm_args: tuple[~typing.Any, ...] | dict[~typing.Any, ~typing.Any] | None = None, activation: type[~torch.nn.modules.module.Module] | None = None, act_args: tuple[~typing.Any, ...] | dict[~typing.Any, ~typing.Any] | None = None, linear_layer: ~collections.abc.Callable[[int, int], ~torch.nn.modules.module.Module] = <class 'torch.nn.modules.linear.Linear'>) list[Module][source]#

Construct a miniblock with given input/output-size, norm layer and activation.

class ModuleWithVectorOutput(output_dim: int)[source]#

Bases: Module

A module that outputs a vector of a known size.

Use from_module to adapt a module to this interface.

Parameters:

output_dim – the dimension of the output vector.

static from_module(module: Module, output_dim: int) ModuleWithVectorOutput[source]#
Parameters:
  • module – the module to adapt.

  • output_dim – dimension of the output vector produced by the module.

get_output_dim() int[source]#
Returns:

the dimension of the output vector.

class ModuleWithVectorOutputAdapter(module: Module, output_dim: int)[source]#

Bases: ModuleWithVectorOutput

Adapts a module with vector output to provide the ModuleWithVectorOutput interface.

Parameters:
  • module – the module to adapt.

  • output_dim – the dimension of the output vector produced by the module.

forward(*args: Any, **kwargs: Any) Any[source]#

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class MLP(*, input_dim: int, output_dim: int = 0, hidden_sizes: ~collections.abc.Sequence[int] = (), norm_layer: type[~torch.nn.modules.module.Module] | ~collections.abc.Sequence[type[~torch.nn.modules.module.Module]] | None = None, norm_args: tuple[~typing.Any, ...] | dict[~typing.Any, ~typing.Any] | ~collections.abc.Sequence[tuple[~typing.Any, ...]] | ~collections.abc.Sequence[dict[~typing.Any, ~typing.Any]] | None = None, activation: type[~torch.nn.modules.module.Module] | ~collections.abc.Sequence[type[~torch.nn.modules.module.Module]] | None = <class 'torch.nn.modules.activation.ReLU'>, act_args: tuple[~typing.Any, ...] | dict[~typing.Any, ~typing.Any] | ~collections.abc.Sequence[tuple[~typing.Any, ...]] | ~collections.abc.Sequence[dict[~typing.Any, ~typing.Any]] | None = None, linear_layer: ~collections.abc.Callable[[int, int], ~torch.nn.modules.module.Module] = <class 'torch.nn.modules.linear.Linear'>, flatten_input: bool = True)[source]#

Bases: ModuleWithVectorOutput

Simple MLP backbone.

Parameters:
  • input_dim – dimension of the input vector.

  • output_dim – dimension of the output vector. If set to 0, there is no explicit final linear layer and the output dimension is the last hidden layer’s dimension.

  • hidden_sizes – shape of MLP passed in as a list, not including input_dim and output_dim.

  • norm_layer – use which normalization before activation, e.g., nn.LayerNorm and nn.BatchNorm1d. Default to no normalization. You can also pass a list of normalization modules with the same length of hidden_sizes, to use different normalization module in different layers. Default to no normalization.

  • activation – which activation to use after each layer, can be both the same activation for all layers if passed in nn.Module, or different activation for different Modules if passed in a list. Default to nn.ReLU.

  • linear_layer – use this module as linear layer. Default to nn.Linear.

  • flatten_input – whether to flatten input data. Default to True.

forward(obs: ndarray | Tensor) Tensor[source]#

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class ActionReprNet(*args, **kwargs)[source]#

Bases: Generic[TRecurrentState], Module, ABC

Abstract base class for neural networks used to compute action-related representations from environment observations, which defines the signature of the forward method.

An action-related representation can be a number of things, including:
  • a distribution over actions in a discrete action space in the form of a vector of unnormalized log probabilities (called “logits” in PyTorch jargon)

  • the Q-values of all actions in a discrete action space

  • the parameters of a distribution (e.g., mean and std. dev. for a Gaussian distribution) over actions in a continuous action space

Initializes internal Module state, shared by both nn.Module and ScriptModule.

abstract forward(obs: Tensor | ndarray | BatchProtocol, state: TRecurrentState | None = None, info: dict[str, Any] | None = None) tuple[Tensor | Sequence[Tensor], TRecurrentState | None][source]#

The main method for tianshou to compute action representations (such as actions, inputs of distributions, Q-values, etc) from env observations. Implementations will always make use of the preprocess_net as the first processing step.

Parameters:
  • obs – the observations from the environment as retrieved from ObsBatchProtocol.obs. If the environment is a dict env, this will be an instance of Batch, otherwise it will be an array (or tensor if your env returns tensors).

  • state – the hidden state of the RNN, if applicable

  • info – the info object from the environment step

Returns:

a tuple (action_repr, hidden_state), where action_repr is either an actual action for the environment or a representation from which it can be retrieved/sampled (e.g., mean and std for a Gaussian distribution), and hidden_state is the new hidden state of the RNN, if applicable.

class ActionReprNetWithVectorOutput(output_dim: int)[source]#

Bases: Generic[T], ActionReprNet[T], ModuleWithVectorOutput

A neural network for the computation of action-related representations which outputs a vector of a known size.

Initializes internal Module state, shared by both nn.Module and ScriptModule.

class Actor(output_dim: int)[source]#

Bases: Generic[T], ActionReprNetWithVectorOutput[T], ABC

Initializes internal Module state, shared by both nn.Module and ScriptModule.

abstract get_preprocess_net() ModuleWithVectorOutput[source]#

Returns the network component that is used for pre-processing, i.e. the component which produces a latent representation, which then is transformed into the final output. This is, therefore, the first part of the network which processes the input. For example, a CNN is often used in Atari examples.

We need this method to be able to share latent representation computations with other networks (e.g. critics) within an algorithm.

Actors that do not have a pre-processing stage can return nn.Identity() (see RandomActor for an example).

class Net(*, state_shape: int | ~collections.abc.Sequence[int], action_shape: ~collections.abc.Sequence[int] | int | ~numpy.int64 = 0, hidden_sizes: ~collections.abc.Sequence[int] = (), norm_layer: type[~torch.nn.modules.module.Module] | ~collections.abc.Sequence[type[~torch.nn.modules.module.Module]] | None = None, norm_args: tuple[~typing.Any, ...] | dict[~typing.Any, ~typing.Any] | ~collections.abc.Sequence[tuple[~typing.Any, ...]] | ~collections.abc.Sequence[dict[~typing.Any, ~typing.Any]] | None = None, activation: type[~torch.nn.modules.module.Module] | ~collections.abc.Sequence[type[~torch.nn.modules.module.Module]] | None = <class 'torch.nn.modules.activation.ReLU'>, act_args: tuple[~typing.Any, ...] | dict[~typing.Any, ~typing.Any] | ~collections.abc.Sequence[tuple[~typing.Any, ...]] | ~collections.abc.Sequence[dict[~typing.Any, ~typing.Any]] | None = None, softmax: bool = False, concat: bool = False, num_atoms: int = 1, dueling_param: tuple[dict[str, ~typing.Any], dict[str, ~typing.Any]] | None = None, linear_layer: ~collections.abc.Callable[[int, int], ~torch.nn.modules.module.Module] = <class 'torch.nn.modules.linear.Linear'>)[source]#

Bases: ActionReprNetWithVectorOutput[Any]

A multi-layer perceptron which outputs an action-related representation.

Parameters:
  • state_shape – int or a sequence of int of the shape of state.

  • action_shape – int or a sequence of int of the shape of action.

  • hidden_sizes – shape of MLP passed in as a list.

  • norm_layer – use which normalization before activation, e.g., nn.LayerNorm and nn.BatchNorm1d. Default to no normalization. You can also pass a list of normalization modules with the same length of hidden_sizes, to use different normalization module in different layers. Default to no normalization.

  • activation – which activation to use after each layer, can be both the same activation for all layers if passed in nn.Module, or different activation for different Modules if passed in a list. Default to nn.ReLU.

  • softmax – whether to apply a softmax layer over the last layer’s output.

  • concat – whether the input shape is concatenated by state_shape and action_shape. If it is True, action_shape is not the output shape, but affects the input shape only.

  • num_atoms – in order to expand to the net of distributional RL. Default to 1 (not use).

  • dueling_param – whether to use dueling network to calculate Q values (for Dueling DQN). If you want to use dueling option, you should pass a tuple of two dict (first for Q and second for V) stating self-defined arguments as stated in class:~tianshou.utils.net.common.MLP. Default to None.

  • linear_layer – use this module constructor, which takes the input and output dimension as input, as linear layer. Default to nn.Linear.

See also

Please refer to MLP for more detailed explanation on the usage of activation, norm_layer, etc.

You can also refer to Actor, Critic, etc, to see how it’s suggested be used.

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(obs: Tensor | ndarray | BatchProtocol, state: T | None = None, info: dict[str, Any] | None = None) tuple[Tensor, T | Any][source]#

Mapping: obs -> flatten (inside MLP)-> logits.

Parameters:
  • obs

  • state – unused and returned as is

  • info – unused

class Recurrent(*, layer_num: int, state_shape: int | Sequence[int], action_shape: Sequence[int] | int | int64, hidden_layer_size: int = 128)[source]#

Bases: ActionReprNetWithVectorOutput[RecurrentStateBatch]

Simple Recurrent network based on LSTM.

Initializes internal Module state, shared by both nn.Module and ScriptModule.

get_preprocess_net() ModuleWithVectorOutput[source]#
forward(obs: Tensor | ndarray | BatchProtocol, state: RecurrentStateBatch | None = None, info: dict[str, Any] | None = None) tuple[Tensor, RecurrentStateBatch][source]#

Mapping: obs -> flatten -> logits.

In the evaluation mode, obs should be with shape [bsz, dim]; in the training mode, obs should be with shape [bsz, len, dim]. See the code and comment for more detail.

Parameters:
  • obs

  • state – either None or a dict with keys ‘hidden’ and ‘cell’

  • info – unused

Returns:

predicted action, next state as dict with keys ‘hidden’ and ‘cell’

class ActorCritic(actor: Module, critic: Module)[source]#

Bases: Module

An actor-critic network for parsing parameters.

Using actor_critic.parameters() instead of set.union or list+list to avoid issue #449.

Parameters:
  • actor (nn.Module) – the actor network.

  • critic (nn.Module) – the critic network.

Initializes internal Module state, shared by both nn.Module and ScriptModule.

class DataParallelNet(net: Module)[source]#

Bases: Module

DataParallel wrapper for training agent with multi-GPU.

This class does only the conversion of input data type, from numpy array to torch’s Tensor. If the input is a nested dictionary, the user should create a similar class to do the same thing.

Parameters:

net – the network to be distributed in different GPUs.

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(obs: Tensor | ndarray | BatchProtocol, *args: Any, **kwargs: Any) tuple[Any, Any][source]#

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class ActionReprNetDataParallelWrapper(net: ActionReprNet)[source]#

Bases: ActionReprNet

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(obs: Tensor | ndarray | BatchProtocol, state: TRecurrentState | None = None, info: dict[str, Any] | None = None) tuple[Tensor, TRecurrentState | None][source]#

The main method for tianshou to compute action representations (such as actions, inputs of distributions, Q-values, etc) from env observations. Implementations will always make use of the preprocess_net as the first processing step.

Parameters:
  • obs – the observations from the environment as retrieved from ObsBatchProtocol.obs. If the environment is a dict env, this will be an instance of Batch, otherwise it will be an array (or tensor if your env returns tensors).

  • state – the hidden state of the RNN, if applicable

  • info – the info object from the environment step

Returns:

a tuple (action_repr, hidden_state), where action_repr is either an actual action for the environment or a representation from which it can be retrieved/sampled (e.g., mean and std for a Gaussian distribution), and hidden_state is the new hidden state of the RNN, if applicable.

class EnsembleLinear(ensemble_size: int, in_feature: int, out_feature: int, bias: bool = True)[source]#

Bases: Module

Linear Layer of Ensemble network.

Parameters:
  • ensemble_size – Number of subnets in the ensemble.

  • in_feature – dimension of the input vector.

  • out_feature – dimension of the output vector.

  • bias – whether to include an additive bias, default to be True.

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(x: Tensor) Tensor[source]#

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class BranchingNet(*, state_shape: int | ~collections.abc.Sequence[int], num_branches: int = 0, action_per_branch: int = 2, common_hidden_sizes: list[int] | None = None, value_hidden_sizes: list[int] | None = None, action_hidden_sizes: list[int] | None = None, norm_layer: type[~torch.nn.modules.module.Module] | None = None, norm_args: tuple[~typing.Any, ...] | dict[~typing.Any, ~typing.Any] | ~collections.abc.Sequence[tuple[~typing.Any, ...]] | ~collections.abc.Sequence[dict[~typing.Any, ~typing.Any]] | None = None, activation: type[~torch.nn.modules.module.Module] | None = <class 'torch.nn.modules.activation.ReLU'>, act_args: tuple[~typing.Any, ...] | dict[~typing.Any, ~typing.Any] | ~collections.abc.Sequence[tuple[~typing.Any, ...]] | ~collections.abc.Sequence[dict[~typing.Any, ~typing.Any]] | None = None)[source]#

Bases: ActionReprNet

Branching dual Q network.

Network for the BranchingDQNPolicy, it uses a common network module, a value module and action “branches” one for each dimension. It allows for a linear scaling of Q-value the output w.r.t. the number of dimensions in the action space.

This network architecture efficiently handles environments with multiple independent action dimensions by using a branching structure. Instead of representing all action combinations (which grows exponentially), it represents each action dimension separately (linear scaling). For example, if there are 3 actions with 3 possible values each, then we would normally need to consider 3^4 = 81 unique actions, whereas with this architecture, we can instead use 3 branches with 4 actions per dimension, resulting in 3 * 4 = 12 values to be considered.

Common use cases include multi-joint robotic control tasks, where each joint can be controlled independently.

For more information, please refer to: arXiv:1711.08946.

Parameters:
  • state_shape – int or a sequence of int of the shape of state.

  • num_branches – number of action dimensions in the environment. Each branch represents one independent action dimension. For example, in a robot with 7 joints, you would set this to 7.

  • action_per_branch – Number of possible discrete values for each action dimension. For example, if each joint can have 3 positions (left, center, right), you would set this to 3.

  • common_hidden_sizes – shape of the common MLP network passed in as a list.

  • value_hidden_sizes – shape of the value MLP network passed in as a list.

  • action_hidden_sizes – shape of the action MLP network passed in as a list.

  • norm_layer – use which normalization before activation, e.g., nn.LayerNorm and nn.BatchNorm1d. Default to no normalization. You can also pass a list of normalization modules with the same length of hidden_sizes, to use different normalization module in different layers. Default to no normalization.

  • activation – which activation to use after each layer, can be both the same activation for all layers if passed in nn.Module, or different activation for different Modules if passed in a list. Default to nn.ReLU.

forward(obs: Tensor | ndarray | BatchProtocol, state: T | None = None, info: dict[str, Any] | None = None) tuple[Tensor, T | None][source]#

Mapping: obs -> model -> logits.

get_dict_state_decorator(state_shape: dict[str, int | Sequence[int]], keys: Sequence[str]) tuple[Callable, int][source]#

A helper function to make Net or equivalent classes (e.g. Actor, Critic) applicable to dict state.

The first return item, decorator_fn, will alter the implementation of forward function of the given class by preprocessing the observation. The preprocessing is basically flatten the observation and concatenate them based on the keys order. The batch dimension is preserved if presented. The result observation shape will be equal to new_state_shape, the second return item.

Parameters:
  • state_shape – A dictionary indicating each state’s shape

  • keys – A list of state’s keys. The flatten observation will be according to this list order.

Returns:

a 2-items tuple decorator_fn and new_state_shape

class AbstractContinuousActorProbabilistic(output_dim: int)[source]#

Bases: Actor, ABC

Type bound for probabilistic actors which output distribution parameters for continuous action spaces.

Initializes internal Module state, shared by both nn.Module and ScriptModule.

class AbstractDiscreteActor(output_dim: int)[source]#

Bases: Actor, ABC

Type bound for discrete actors.

For on-policy algos like Reinforce, this typically directly outputs unnormalized log probabilities, which can be interpreted as “logits” in conjunction with a torch.distributions.Categorical instance.

In Tianshou, discrete actors are also used for computing action distributions within Q-learning type algorithms (e.g., DQN). In this case, the observations are mapped to a vector of Q-values (one for each action). In other words, the component is actually a critic, not an actor in the traditional sense. Note that when sampling actions, the Q-values can be interpreted as inputs for a torch.distributions.Categorical instance, similar to the on-policy case mentioned above.

Initializes internal Module state, shared by both nn.Module and ScriptModule.

class RandomActor(action_space: Box | Discrete)[source]#

Bases: AbstractContinuousActorProbabilistic, AbstractDiscreteActor

An actor that returns random actions.

For continuous action spaces, forward returns a batch of random actions sampled from the action space. For discrete action spaces, forward returns a batch of n-dimensional arrays corresponding to the uniform distribution over the n possible actions (same interface as in Actor).

Initializes internal Module state, shared by both nn.Module and ScriptModule.

property action_space: Box | Discrete#
property space_info: ActionSpaceInfo#
get_preprocess_net() ModuleWithVectorOutput[source]#

Returns the network component that is used for pre-processing, i.e. the component which produces a latent representation, which then is transformed into the final output. This is, therefore, the first part of the network which processes the input. For example, a CNN is often used in Atari examples.

We need this method to be able to share latent representation computations with other networks (e.g. critics) within an algorithm.

Actors that do not have a pre-processing stage can return nn.Identity() (see RandomActor for an example).

get_output_dim() int[source]#
Returns:

the dimension of the output vector.

property is_discrete: bool#
forward(obs: Tensor | ndarray | BatchProtocol, state: T | None = None, info: dict[str, Any] | None = None) tuple[Tensor, T | None][source]#

The main method for tianshou to compute action representations (such as actions, inputs of distributions, Q-values, etc) from env observations. Implementations will always make use of the preprocess_net as the first processing step.

Parameters:
  • obs – the observations from the environment as retrieved from ObsBatchProtocol.obs. If the environment is a dict env, this will be an instance of Batch, otherwise it will be an array (or tensor if your env returns tensors).

  • state – the hidden state of the RNN, if applicable

  • info – the info object from the environment step

Returns:

a tuple (action_repr, hidden_state), where action_repr is either an actual action for the environment or a representation from which it can be retrieved/sampled (e.g., mean and std for a Gaussian distribution), and hidden_state is the new hidden state of the RNN, if applicable.

compute_action_batch(obs: Tensor | ndarray | BatchProtocol) Tensor[source]#