lit_sb3.modules

source package lit_sb3.modules

Classes

• BaseFeaturesExtractor — Base class that represents a features extractor.

• BatchRenorm — BatchRenorm Module (https://arxiv.org/abs/1702.03275).

• BatchRenorm1d

• CombinedExtractor — Combined features extractor for Dict observation spaces.

• FlattenExtractor — Feature extract that flatten the input.

• MlpExtractor — Constructs an MLP that receives the output from a previous features extractor (i.e. a CNN) or directly.

• NatureCNN — CNN from DQN Nature paper.

Functions

• create_mlp — Create a multi layer perceptron (MLP), which is.

• get_actor_critic_arch — Get the actor and critic network architectures for off-policy actor-critic algorithms (SAC, TD3, DDPG).

source class BaseFeaturesExtractor(observation_space: gym.Space, features_dim: int = 0)

Bases : nn.Module

Base class that represents a features extractor.

Parameters

• observation_space : gym.Space — The observation space of the environment

• features_dim : int — Number of features extracted.

Attributes

• features_dim : int — The number of features that the extractor outputs.

source property BaseFeaturesExtractor.features_dim: int

The number of features that the extractor outputs.

source class BatchRenorm(num_features: int, eps: float = 0.001, momentum: float = 0.01, affine: bool = True, warmup_steps: int = 100000)

Bases : torch.nn.Module

BatchRenorm Module (https://arxiv.org/abs/1702.03275).

Adapted to Pytorch from https://github.com/araffin/sbx/blob/master/sbx/common/jax_layers.py.

BatchRenorm is an improved version of vanilla BatchNorm. Contrary to BatchNorm, BatchRenorm uses the running statistics for normalizing the batches after a warmup phase. This makes it less prone to suffer from "outlier" batches that can happen during very long training runs and, therefore, is more robust during long training runs.

During the warmup phase, it behaves exactly like a BatchNorm layer. After the warmup phase, the running statistics are used for normalization. The running statistics are updated during training mode. During evaluation mode, the running statistics are used for normalization but not updated.

Parameters

• num_features : int — Number of features in the input tensor.

• eps : float — A value added to the variance for numerical stability.

• momentum : float — The value used for the ra_mean and ra_var (running average) computation. It controls the rate of convergence for the batch renormalization statistics.

• affine : bool — A boolean value that when set to True, this module has learnable affine parameters. Default: True

• warmup_steps : int — Number of warum steps that are performed before the running statistics are used for normalization. During the warump phase, the batch statistics are used.

Methods

• forward — Normalize the input tensor.

• extra_repr

source method BatchRenorm.forward(x: torch.Tensor) → torch.Tensor

Normalize the input tensor.

Parameters

• x : torch.Tensor — Input tensor

Returns

• torch.Tensor — Normalized tensor.

source method BatchRenorm.extra_repr() → str

source class BatchRenorm1d(num_features: int, eps: float = 0.001, momentum: float = 0.01, affine: bool = True, warmup_steps: int = 100000)

Bases : BatchRenorm

source class CombinedExtractor(observation_space: spaces.Dict, cnn_output_dim: int = 256, normalized_image: bool = False)

Bases : BaseFeaturesExtractor

Combined features extractor for Dict observation spaces.

Builds a features extractor for each key of the space. Input from each space is fed through a separate submodule (CNN or MLP, depending on input shape), the output features are concatenated and fed through additional MLP network ("combined").

Parameters

• observation_space : spaces.Dict

• cnn_output_dim : int — Number of features to output from each CNN submodule(s). Defaults to 256 to avoid exploding network sizes.

• normalized_image : bool — Whether to assume that the image is already normalized or not (this disables dtype and bounds checks): when True, it only checks that the space is a Box and has 3 dimensions. Otherwise, it checks that it has expected dtype (uint8) and bounds (values in [0, 255]).

Attributes

• features_dim : int — The number of features that the extractor outputs.

Methods

• forward

source method CombinedExtractor.forward(observations: TensorDict) → torch.Tensor

source class FlattenExtractor(observation_space: gym.Space)

Bases : BaseFeaturesExtractor

Feature extract that flatten the input.

Used as a placeholder when feature extraction is not needed.

Parameters

• observation_space : gym.Space — The observation space of the environment

Attributes

• features_dim : int — The number of features that the extractor outputs.

Methods

• forward

source method FlattenExtractor.forward(observations: torch.Tensor) → torch.Tensor

source class MlpExtractor(feature_dim: int, net_arch: list[int] | dict[str, list[int]], activation_fn: type[nn.Module], device: torch.device | str = 'auto')

Bases : nn.Module

Constructs an MLP that receives the output from a previous features extractor (i.e. a CNN) or directly.

the observations (if no features extractor is applied) as an input and outputs a latent representation for the policy and a value network.

The net_arch parameter allows to specify the amount and size of the hidden layers. It can be in either of the following forms: 1. dict(vf=[<list of layer sizes>], pi=[<list of layer sizes>]): to specify the amount and size of the layers in the policy and value nets individually. If it is missing any of the keys (pi or vf), zero layers will be considered for that key. 2. [<list of layer sizes>]: "shortcut" in case the amount and size of the layers in the policy and value nets are the same. Same as dict(vf=int_list, pi=int_list) where int_list is the same for the actor and critic.

.. note:: If a key is not specified or an empty list is passed [], a linear network will be used.

Parameters

• feature_dim : int — Dimension of the feature vector (can be the output of a CNN)

• net_arch : list[int] | dict[str, list[int]] — The specification of the policy and value networks. See above for details on its formatting.

• activation_fn : type[nn.Module] — The activation function to use for the networks.

• device : torch.device | str — PyTorch device.

Methods

• forward — Returns.

• forward_actor

• forward_critic

source method MlpExtractor.forward(features: torch.Tensor) → tuple[torch.Tensor, torch.Tensor]

Returns.

latent_policy, latent_value of the specified network. If all layers are shared, then latent_policy == latent_value.

source method MlpExtractor.forward_actor(features: torch.Tensor) → torch.Tensor

source method MlpExtractor.forward_critic(features: torch.Tensor) → torch.Tensor

source class NatureCNN(observation_space: gym.Space, features_dim: int = 512, normalized_image: bool = False)

Bases : BaseFeaturesExtractor

CNN from DQN Nature paper.

Mnih, Volodymyr, et al.
"Human-level control through deep reinforcement learning."
Nature 518.7540 (2015): 529-533.

Parameters

• observation_space : gym.Space — The observation space of the environment

• features_dim : int — Number of features extracted. This corresponds to the number of unit for the last layer.

• normalized_image : bool — Whether to assume that the image is already normalized or not (this disables dtype and bounds checks): when True, it only checks that the space is a Box and has 3 dimensions. Otherwise, it checks that it has expected dtype (uint8) and bounds (values in [0, 255]).

Attributes

• features_dim : int — The number of features that the extractor outputs.

Methods

• forward

source method NatureCNN.forward(observations: torch.Tensor) → torch.Tensor

source create_mlp(input_dim: int, output_dim: int, net_arch: list[int], activation_fn: type[nn.Module] = nn.ReLU, squash_output: bool = False, with_bias: bool = True, pre_linear_modules: list[type[nn.Module]] | None = None, post_linear_modules: list[type[nn.Module]] | None = None) → list[nn.Module]

Create a multi layer perceptron (MLP), which is.

a collection of fully-connected layers each followed by an activation function.

Parameters

• input_dim : int — Dimension of the input vector

• output_dim : int — Dimension of the output (last layer, for instance, the number of actions)

• net_arch : list[int] — Architecture of the neural net It represents the number of units per layer. The length of this list is the number of layers.

• activation_fn : type[nn.Module] — The activation function to use after each layer.

• squash_output : bool — Whether to squash the output using a Tanh activation function

• with_bias : bool — If set to False, the layers will not learn an additive bias

• pre_linear_modules : list[type[nn.Module]] | None — List of nn.Module to add before the linear layers. These modules should maintain the input tensor dimension (e.g. BatchNorm). The number of input features is passed to the module's constructor. Compared to post_linear_modules, they are used before the output layer (output_dim > 0).

• post_linear_modules : list[type[nn.Module]] | None — List of nn.Module to add after the linear layers (and before the activation function). These modules should maintain the input tensor dimension (e.g. Dropout, LayerNorm). They are not used after the output layer (output_dim > 0). The number of input features is passed to the module's constructor.

Returns

• list[nn.Module] — The list of layers of the neural network

source get_actor_critic_arch(net_arch: list[int] | dict[str, list[int]]) → tuple[list[int], list[int]]

Get the actor and critic network architectures for off-policy actor-critic algorithms (SAC, TD3, DDPG).

The net_arch parameter allows to specify the amount and size of the hidden layers, which can be different for the actor and the critic. It is assumed to be a list of ints or a dict.

1. If it is a list, actor and critic networks will have the same architecture. The architecture is represented by a list of integers (of arbitrary length (zero allowed)) each specifying the number of units per layer. If the number of ints is zero, the network will be linear.
2. If it is a dict, it should have the following structure: dict(qf=[<critic network architecture>], pi=[<actor network architecture>]). where the network architecture is a list as described in 1.

For example, to have actor and critic that share the same network architecture, you only need to specify net_arch=[256, 256] (here, two hidden layers of 256 units each).

If you want a different architecture for the actor and the critic, then you can specify net_arch=dict(qf=[400, 300], pi=[64, 64]).

.. note:: Compared to their on-policy counterparts, no shared layers (other than the features extractor) between the actor and the critic are allowed (to prevent issues with target networks).

Parameters

• net_arch : list[int] | dict[str, list[int]] — The specification of the actor and critic networks. See above for details on its formatting.

Returns

• tuple[list[int], list[int]] — The network architectures for the actor and the critic