From 7c9e75030b54020ff7fe79b6533580b9983aa04f Mon Sep 17 00:00:00 2001
From: VedantDave <vedant.dave@unileoben.ac.at>
Date: Fri, 31 Mar 2023 19:12:46 +0200
Subject: [PATCH] Updating value model to be stochastic

---
 DPI/models.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/DPI/models.py b/DPI/models.py
index ed337a2..f2d1d4c 100644
--- 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):