Modules

MLP

RNN

CNN

class rsl_rl.modules.cnn.CNN[source]

Convolutional Neural Network.

The CNN network is a sequence of convolutional layers, optional normalization layers, optional activation functions, and optional pooling. The final output can be flattened.

__init__(input_dim, input_channels, output_channels, kernel_size, stride=1, dilation=1, padding='none', norm='none', activation='elu', max_pool=False, global_pool='none', flatten=True)[source]

Initialize the CNN.

Parameters:
  • input_dim (tuple[int, int]) – Height and width of the input.

  • input_channels (int) – Number of input channels.

  • output_channels (tuple[int, ...] | list[int]) – List of output channels for each convolutional layer.

  • kernel_size (int | tuple[int, ...] | list[int]) – List of kernel sizes for each convolutional layer or a single kernel size for all layers.

  • stride (int | tuple[int, ...] | list[int]) – List of strides for each convolutional layer or a single stride for all layers.

  • dilation (int | tuple[int, ...] | list[int]) – List of dilations for each convolutional layer or a single dilation for all layers.

  • padding (str) – Padding type to use. Either ‘none’, ‘zeros’, ‘reflect’, ‘replicate’, or ‘circular’.

  • norm (str | tuple[str] | list[str]) – List of normalization types for each convolutional layer or a single type for all layers. Either ‘none’, ‘batch’, or ‘layer’.

  • activation (str) – Activation function to use.

  • max_pool (bool | tuple[bool] | list[bool]) – List of booleans indicating whether to apply max pooling after each convolutional layer or a single boolean for all layers.

  • global_pool (str) – Global pooling type to apply at the end. Either ‘none’, ‘max’, or ‘avg’.

  • flatten (bool) – Whether to flatten the output tensor.

Return type:

None

property output_channels: int | None

Get the number of output channels or None if output is flattened.

property output_dim: tuple[int, int] | int

Get the output height and width or total output dimension if output is flattened.

init_weights()[source]

Initialize the weights of the CNN with Kaiming initialization.

Return type:

None

forward(x)[source]

Forward pass of the CNN.

Parameters:

x (torch.Tensor)

Return type:

torch.Tensor

Normalization

Distribution