add log_prob methods to ConstantDelay and DoublePoissonDelay for vectorized log probability lookups.

Author: yichao-liang

predicators/approaches/pp_param_learning_approach.py

@@ -333,11 +333,34 @@ def elbo_torch(
         for starts, gp in zip(start_times, ground_processes):
             if len(starts) > 0:
                 s0 = starts[0]
-                for t in range(s0 + 1, num_time_steps):
-                    delay_prob = gp.delay_distribution.probability(t - s0)
-                    if delay_prob > 1e-9:
-                        ll = ll + guide[gp][t] * torch.log(
-                            torch.tensor(delay_prob))
+                # Relevant time steps for this gp's delay calculation: s0+1 to 
+                # num_time_steps-1
+                # These correspond to delay values: 1 to num_time_steps-1-s0
+                # Check if there are any time steps for delay
+                if s0 + 1 < num_time_steps:
+                    # Create a tensor of delay values that occurred
+                    delay_values = torch.arange(1,
+                                                num_time_steps - s0,
+                                                dtype=torch.long)
+                    # Corresponding time indices for guide probabilities
+                    t_indices_for_guide = torch.arange(s0 + 1,
+                                                       num_time_steps,
+                                                       dtype=torch.long)
+
+                    # Get log prob for all possible delay values at once
+                    all_delay_log_probs = gp.delay_distribution.log_prob(
+                        delay_values)
+
+                    # Get the slice of guide prob relevant to these time steps
+                    guide_slice = guide[gp][t_indices_for_guide]
+
+                    # Mask for valid log prob (not -inf)
+                    valid_mask = ~torch.isneginf(all_delay_log_probs)
+
+                    # Add to ll only for terms where log probability is valid
+                    if valid_mask.any():
+                        ll = ll + torch.sum(guide_slice[valid_mask] *
+                                            all_delay_log_probs[valid_mask])
 
         # -----------------------------------------------------------------
         # 3.  Entropy of the variational distributions

predicators/ground_truth_models/boil/processes.py

@@ -8,8 +8,8 @@
 from predicators.structs import CausalProcess, EndogenousProcess, \
     ExogenousProcess, LiftedAtom, ParameterizedOption, Predicate, Type, \
     Variable
-from predicators.utils import CMPDelay, ConstantDelay, GaussianDelay, \
-    null_sampler, DoublePoissonDelay
+from predicators.utils import CMPDelay, ConstantDelay, DoublePoissonDelay, \
+    GaussianDelay, null_sampler
 
 
 class PyBulletBoilGroundTruthProcessFactory(GroundTruthProcessFactory):
@@ -413,7 +413,8 @@ def get_processes(
             elif CFG.boil_use_normal_delay:
                 delay_distribution = GaussianDelay(mean=1, std=0.2, rng=rng)
             elif CFG.boil_use_cmp_delay:
-                delay_distribution = CMPDelay(1, 1) # Assumed from other mean=1 cases
+                delay_distribution = CMPDelay(
+                    1, 1)  # Assumed from other mean=1 cases
             else:
                 delay_distribution = DoublePoissonDelay(mu=1, phi=50)
             declare_complete_process = EndogenousProcess(
@@ -452,7 +453,8 @@ def get_processes(
         elif CFG.boil_use_cmp_delay:
             delay_distribution = CMPDelay(100, 2.9)
         else:
-            delay_distribution = DoublePoissonDelay(mu=4, phi=50) # mu=4 from ConstantDelay(4)
+            delay_distribution = DoublePoissonDelay(
+                mu=4, phi=50)  # mu=4 from ConstantDelay(4)
         fill_jug_process = ExogenousProcess("FillJug", parameters,
                                             condition_at_start,
                                             condition_overall, set(),
@@ -572,7 +574,8 @@ def get_processes(
             elif CFG.boil_use_normal_delay:
                 delay_distribution = GaussianDelay(mean=3, std=0.2, rng=rng)
             elif CFG.boil_use_cmp_delay:
-                delay_distribution = CMPDelay(55, 3) # Assumed from other mean=3 cases
+                delay_distribution = CMPDelay(
+                    55, 3)  # Assumed from other mean=3 cases
             else:
                 delay_distribution = DoublePoissonDelay(mu=3, phi=50)
             human_happy_process = ExogenousProcess("HumanHappy", parameters,

predicators/utils.py

@@ -4311,6 +4311,10 @@ class ConstantDelay(DelayDistribution):
 
     def __init__(self, delay: int):
         self.delay = delay
+        self._log_pmf_value = torch.tensor(0.0,
+                                           dtype=torch.get_default_dtype())
+        self._neg_inf = torch.tensor(float('-inf'),
+                                     dtype=torch.get_default_dtype())
 
     def sample(self):
         return self.delay
@@ -4326,9 +4330,24 @@ def set_parameters(self, parameters):
     def probability(self, k: int) -> float:
         return 1.0 if k == self.delay else 0.0
 
+    def log_prob(self, k: Union[int, torch.Tensor]) -> torch.Tensor:
+        """Fast vectorized lookup of log P(delay = k)."""
+        if not isinstance(k, torch.Tensor):
+            k_tensor = torch.tensor(k, dtype=torch.long)
+        else:
+            k_tensor = k.long()
+
+        # Initialize with -inf
+        log_probs = torch.full_like(k_tensor,
+                                    self._neg_inf,
+                                    dtype=torch.get_default_dtype())
+        # Set log(1.0) = 0.0 where k == self.delay
+        log_probs[k_tensor == self.delay] = self._log_pmf_value
+        return log_probs
+
     @cached_property
     def _str(self) -> str:
-        return f"ConstantDelay({self.delay})"
+        return f"ConstantDelay({self.delay:.4f})"
 
 
 class GaussianDelay(DelayDistribution):
@@ -4438,15 +4457,22 @@ def __init__(self, mu: float, phi: float, max_k: int = 50):
     def _update_cache(self) -> None:
         ks = torch.arange(self._max_k, dtype=torch.get_default_dtype())
         log_fact = gammaln(ks + 1)
-        log_C = -0.5 * (torch.log(torch.tensor(2 * torch.pi)) +
-                        torch.log(self.phi) + torch.log(self.mu))
-        log_p = (log_C + self.phi * (ks * torch.log(self.mu) - log_fact) -
-                 self.phi * self.mu)
-        log_p -= torch.max(log_p)
-        pmf = torch.exp(log_p)
-        pmf /= pmf.sum()
-        self._pmf = pmf
-        self._cdf = torch.cumsum(pmf, dim=0)
+        log_C = -0.5 * (
+            torch.log(torch.tensor(2 * torch.pi, dtype=self.mu.dtype))
+            +  # Ensure tensor has same dtype
+            torch.log(self.phi) + torch.log(self.mu))
+        # Calculate unnormalized log probabilities
+        log_p_unnormalized = (log_C + self.phi *
+                              (ks * torch.log(self.mu) - log_fact) -
+                              self.phi * self.mu)
+
+        # Calculate normalized log PMF using logsumexp for numerical stability
+        log_normalization_constant = torch.logsumexp(log_p_unnormalized, dim=0)
+        self._log_pmf = log_p_unnormalized - log_normalization_constant
+
+        # Calculate PMF from the normalized log PMF
+        self._pmf = torch.exp(self._log_pmf)
+        self._cdf = torch.cumsum(self._pmf, dim=0)
 
     def set_parameters(self, parameters: Sequence[float]) -> None:
         """Update μ and φ, then rebuild the cache."""
@@ -4465,6 +4491,32 @@ def probability(self, k: int) -> float:
             return float(self._pmf[k])
         return 0.0
 
+    def log_prob(self, k: Union[int, torch.Tensor]) -> torch.Tensor:
+        """Fast vectorized lookup of log P(delay = k)."""
+        if not isinstance(k, torch.Tensor):
+            k_tensor = torch.tensor(
+                k, dtype=torch.long)  # Use long for potential indexing
+        else:
+            k_tensor = k.long()  # Ensure long type for indexing
+
+        original_shape = k_tensor.shape
+        k_flat = k_tensor.flatten()
+
+        # Initialize log_probs with -inf, matching the default dtype
+        log_probs_flat = torch.full_like(k_flat,
+                                         float('-inf'),
+                                         dtype=torch.get_default_dtype())
+
+        assert self._log_pmf is not None, "Cache not updated, _log_pmf is None."
+
+        valid_mask = (k_flat >= 0) & (k_flat < self._max_k)
+        valid_indices_in_k = k_flat[valid_mask]
+
+        if valid_indices_in_k.numel() > 0:
+            log_probs_flat[valid_mask] = self._log_pmf[valid_indices_in_k]
+
+        return log_probs_flat.reshape(original_shape)
+
     def sample(self) -> int:
         u = torch.rand(1).item()
         return int(torch.searchsorted(self._cdf, torch.tensor(u)))