from torch import nn
import math
[docs]class BaseNNEstimator(nn.Module):
def __init__(
self,
observation_shape: tuple,
num_hidden_layers: int = 0,
num_outputs: int = 1,
in_features: int = 64,
out_features: int = 64,
use_bias: bool = True,
):
super().__init__()
num_input_features = math.prod(observation_shape)
[docs] self.num_hidden_layers = num_hidden_layers
[docs] self.in_features = in_features
[docs] self.out_features = out_features
[docs] self.hidden_layers = nn.ModuleList()
if self.num_hidden_layers <= 0:
self.linear_out = nn.Linear(
in_features=num_input_features, out_features=num_outputs, bias=use_bias
)
else:
self.hidden_layers.append(
nn.Sequential(
nn.Linear(
in_features=num_input_features,
out_features=self.out_features,
bias=use_bias,
),
nn.ReLU(),
)
)
for _ in range(self.num_hidden_layers - 1):
self.hidden_layers.append(self._create_hidden_layer(use_bias=use_bias))
self.linear_out = nn.Linear(
in_features=self.in_features, out_features=num_outputs, bias=use_bias
)
[docs] def _create_hidden_layer(self, use_bias: bool):
return nn.Sequential(
nn.Linear(self.in_features, self.out_features, bias=use_bias), nn.ReLU()
)
[docs] def forward(self, x):
x = x.view(-1) # Flatten input tensor
for hidden_layer in self.hidden_layers:
x = hidden_layer(x)
out = self.linear_out(x)
return out