Updating value model to be stochastic

2023-03-31 19:12:46 +02:00 · 2023-03-31 19:12:46 +02:00 · 7c9e75030b
commit 7c9e75030b
parent e6e11f90b7
1 changed files with 2 additions and 1 deletions
--- a/DPI/models.py
+++ b/DPI/models.py
@ -149,7 +149,8 @@ class ValueModel(nn.Module):

    def forward(self, state):
        value = self.value_model(state)
-        return value
+        value_dist = torch.distributions.independent.Independent(torch.distributions.Normal(value, 1), 1)
+        return value_dist
    
        
 class TransitionModel(nn.Module):