diff --git a/DPI/models.py b/DPI/models.py index e6adcf4..ed337a2 100644 --- a/DPI/models.py +++ b/DPI/models.py @@ -132,6 +132,25 @@ class ActionDecoder(nn.Module): return sample +class ValueModel(nn.Module): + def __init__(self, state_size, hidden_size, num_layers=4): + super().__init__() + self.state_size = state_size + self.hidden_size = hidden_size + self.num_layers = num_layers + + layers = [] + for i in range(self.num_layers): + input_channels = state_size if i == 0 else self.hidden_size + output_channels = self.hidden_size if i!= self.num_layers-1 else 1 + layers.append(nn.Linear(input_channels, output_channels)) + layers.append(nn.ReLU()) + self.value_model = nn.Sequential(*layers) + + def forward(self, state): + value = self.value_model(state) + return value + class TransitionModel(nn.Module): def __init__(self, state_size, hidden_size, action_size, history_size):