Merge updated files

src/safe_options/policy.py

@@ -0,0 +1,35 @@
+import torch
+import torch.nn as nn
+from torch.distributions import Categorical
+from torch.distributions.kl import kl_divergence
+from core.policy import SetDiscretePolicy
+
+class SetMaskedDiscretePolicy(SetDiscretePolicy):
+
+    def forward(self, observation, safe_actions):
+        return torch.cat((super().forward(observation), safe_actions), -1)
+    
+    def torch_dist(self, dist):
+        logits = dist[..., :self.action_dim]
+        z = dist[..., self.action_dim:]
+        a = super().torch_dist(logits).probs
+        return Categorical(probs=a*z)
+    
+    def unsafe_probability_mass(self, dist):
+        logits = dist[..., :self.action_dim]
+        z = dist[..., self.action_dim:]
+        a = super().torch_dist(logits).probs
+        return (a * (1 - z)).sum(-1)
+    
+    # def torch_dist_nomask(self, dist):
+    #     print('no mask logprob')
+    #     logits = dist[..., :self.action_dim]
+    #     return super().torch_dist(logits)
+
+    # def log_prob(self, dist, actions):
+    #     return self.torch_dist_nomask(dist).log_prob(actions)
+    
+    # def kl_divergence(self, dist1, dist2):
+    #     d1 = self.torch_dist_nomask(dist1)
+    #     d2 = self.torch_dist_nomask(dist2)
+    #     return kl_divergence(d1, d2)