Adding value model

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):