algorithms/common/models/mlp.py

from typing import Type, Optional

import torch
from torch import nn as nn


class SimpleMlp(nn.Module):
    """
    A class for very simple multi layer perceptron
    """

    def __init__(
        self,
        in_dim=2,
        out_dim=1,
        hidden_dim=64,
        n_layers=2,
        activation: Type[nn.Module] = nn.ReLU,
        output_activation: Optional[Type[nn.Module]] = None,
    ):
        super(SimpleMlp, self).__init__()
        layers = [nn.Linear(in_dim, hidden_dim), activation()]
        layers.extend(
            [nn.Linear(hidden_dim, hidden_dim), activation()] * (n_layers - 2)
        )
        layers.append(nn.Linear(hidden_dim, out_dim))
        if output_activation:
            layers.append(output_activation())
        self.net = nn.Sequential(*layers)

    def forward(self, x):
        return self.net(x)