Merge branch 'XanderJC:main' into main

vanderschaarlab · Oct 2, 2024 · 7764b08 · 7764b08
2 parents a8bed6e + 13b0690
commit 7764b08
Show file tree

Hide file tree

Showing 7 changed files with 5 additions and 5 deletions.
diff --git a/medkit/domains/icu.py b/medkit/domains/icu.py
@@ -82,10 +82,10 @@ def __init__(self, y_dim=2):
         VAE_config = {
             "latent_size": 10,
             "hidden_units": 100,
-            "lr": 1e-4,
+            "lr": 1e-5,
             "hidden_layers": 3,
             "adam_betas": (0.9, 0.9),
-            "epochs": 10,
+            "epochs": 200,
         }
 
         self.env_config_dict = {

diff --git a/medkit/domains/ward.py b/medkit/domains/ward.py
@@ -86,7 +86,7 @@ def __init__(self, y_dim=2):
             "lr": 1e-4,
             "hidden_layers": 3,
             "adam_betas": (0.9, 0.9),
-            "epochs": 10,
+            "epochs": 200,
         }
 
         self.env_config_dict = {

diff --git a/medkit/environments/SequentialVAE.py b/medkit/environments/SequentialVAE.py
@@ -198,7 +198,7 @@ def step(self, action):
         action = action.reshape((1, 1))
         action_one_hot = F.one_hot(action, self.domain.y_dim)
 
-        x = torch.cat((self.prev_latent, action_one_hot, self.static), 2)
+        x = torch.cat((self.prev_obs, action_one_hot, self.static.expand(1, 1, -1)), 2)
 
         out, (self.hn, self.cn) = self.model.lstm(x, (self.hn, self.cn))
 

diff --git a/medkit/environments/TForce.py b/medkit/environments/TForce.py
@@ -102,7 +102,7 @@ def step(self, action):
         action = action.reshape((1, 1))
         action_one_hot = F.one_hot(action, self.domain.y_dim)
 
-        x = torch.cat((self.prev_obs, action_one_hot, self.static), 2)
+        x = torch.cat((self.prev_obs, action_one_hot, self.static.expand(1, 1, -1)), 2)
 
         out, (self.hn, self.cn) = self.model.lstm(x, (self.hn, self.cn))
 

diff --git a/medkit/environments/saved_models/icu_8_statespace.pth b/medkit/environments/saved_models/icu_8_statespace.pth
diff --git a/medkit/initialisers/saved_models/icu_VAE.pth b/medkit/initialisers/saved_models/icu_VAE.pth
diff --git a/medkit/initialisers/saved_models/ward_VAE.pth b/medkit/initialisers/saved_models/ward_VAE.pth