Support log_prob(value, summed) for DiagMultivariateNormal

HorizonRobotics · Dec 16, 2024 · 230c47f · 230c47f
1 parent 3530d14
commit 230c47f
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/alf/utils/dist_utils.py b/alf/utils/dist_utils.py
@@ -373,6 +373,12 @@ def __init__(self, loc, scale):
     def stddev(self):
         return self.base_dist.stddev
 
+    def log_prob(self, value, summed=True):
+        if summed:
+            return self.base_dist.log_prob(value).sum(-1)
+        else:
+            return self.base_dist.log_prob(value)
+
 
 @alf.configurable(whitelist=['eps'])
 class Beta(td.Beta):