benchmarl.models.Mlp

class Mlp(*args, **kwargs)[source]

Bases: Model

Multi layer perceptron model.

Parameters:

  • num_cells (int or Sequence[int], optional) – number of cells of every layer in between the input and output. If an integer is provided, every layer will have the same number of cells. If an iterable is provided, the linear layers out_features will match the content of num_cells.

  • layer_class (Type[nn.Module]) – class to be used for the linear layers;

  • activation_class (Type[nn.Module]) – activation class to be used.

  • activation_kwargs (dict, optional) – kwargs to be used with the activation class;

  • norm_class (Type, optional) – normalization class, if any.

  • norm_kwargs (dict, optional) – kwargs to be used with the normalization layers;

  • num_feature_dims – number of dimensions to be considered as features.

_forward(tensordict: TensorDictBase) → TensorDictBase[source]

Method to implement for the forward pass of the model. It should read self.in_keys, process it and write self.out_key.

Parameters:

tensordict (TensorDictBase) – the input td

Returns: the input td with the written self.out_key