pyRDDLGym-jax 2.2__tar.gz → 2.3__tar.gz

{pyrddlgym_jax-2.2 → pyrddlgym_jax-2.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: pyRDDLGym-jax
-Version: 2.2
+Version: 2.3
 Summary: pyRDDLGym-jax: automatic differentiation for solving sequential planning problems in JAX.
 Home-page: https://github.com/pyrddlgym-project/pyRDDLGym-jax
 Author: Michael Gimelfarb, Ayal Taitler, Scott Sanner

pyrddlgym_jax-2.3/pyRDDLGym_jax/__init__.py

@@ -0,0 +1 @@
+__version__ = '2.3'

{pyrddlgym_jax-2.2 → pyrddlgym_jax-2.3}/pyRDDLGym_jax/core/compiler.py

@@ -1019,6 +1019,9 @@ class JaxRDDLCompiler:
     # UnnormDiscrete: complete (subclass uses Gumbel-softmax)
     # Discrete(p): complete (subclass uses Gumbel-softmax)
     # UnnormDiscrete(p): complete (subclass uses Gumbel-softmax)
+    # Poisson (subclass uses Gumbel-softmax or Poisson process trick)
+    # Binomial (subclass uses Gumbel-softmax or Normal approximation)
+    # NegativeBinomial (subclass uses Poisson-Gamma mixture)
     
     # distributions which seem to support backpropagation (need more testing):
     # Beta
@@ -1026,11 +1029,8 @@ class JaxRDDLCompiler:
     # Gamma
     # ChiSquare   
     # Dirichlet
-    # Poisson (subclass uses Gumbel-softmax or Poisson process trick)
 
     # distributions with incomplete reparameterization support (TODO):
-    # Binomial
-    # NegativeBinomial
     # Multinomial
     
     def _jax_random(self, expr, init_params):
@@ -1299,8 +1299,17 @@ class JaxRDDLCompiler:
     def _jax_negative_binomial(self, expr, init_params):
         ERR = JaxRDDLCompiler.ERROR_CODES['INVALID_PARAM_NEGATIVE_BINOMIAL']
         JaxRDDLCompiler._check_num_args(expr, 2)
-        
         arg_trials, arg_prob = expr.args
+
+        # if prob is non-fluent, always use the exact operation
+        if self.compile_non_fluent_exact \
+        and not self.traced.cached_is_fluent(arg_trials) \
+        and not self.traced.cached_is_fluent(arg_prob):
+            negbin_op = self.EXACT_OPS['sampling']['NegativeBinomial']
+        else:
+            negbin_op = self.OPS['sampling']['NegativeBinomial']
+        jax_op = negbin_op(expr.id, init_params)
+
         jax_trials = self._jax(arg_trials, init_params)
         jax_prob = self._jax(arg_prob, init_params)
         
@@ -1308,11 +1317,8 @@ class JaxRDDLCompiler:
         def _jax_wrapped_distribution_negative_binomial(x, params, key):
             trials, key, err2, params = jax_trials(x, params, key)       
             prob, key, err1, params = jax_prob(x, params, key)
-            trials = jnp.asarray(trials, dtype=self.REAL)
-            prob = jnp.asarray(prob, dtype=self.REAL)
             key, subkey = random.split(key)
-            dist = tfp.distributions.NegativeBinomial(total_count=trials, probs=prob)
-            sample = jnp.asarray(dist.sample(seed=subkey), dtype=self.INT)
+            sample, params = jax_op(subkey, trials, prob, params)
             out_of_bounds = jnp.logical_not(jnp.all(
                 (prob >= 0) & (prob <= 1) & (trials > 0)))
             err = err1 | err2 | (out_of_bounds * ERR)

{pyrddlgym_jax-2.2 → pyrddlgym_jax-2.3}/pyRDDLGym_jax/core/logic.py

@@ -29,15 +29,27 @@
 #
 # ***********************************************************************
 
-from typing import Callable, Dict, Union
+import traceback
+from typing import Callable, Dict, Tuple, Union
 
 import jax
 import jax.numpy as jnp
 import jax.random as random
 import jax.scipy as scipy 
 
+from pyRDDLGym.core.debug.exception import raise_warning
 
-def enumerate_literals(shape, axis, dtype=jnp.int32):
+# more robust approach - if user does not have this or broken try to continue
+try:
+    from tensorflow_probability.substrates import jax as tfp
+except Exception:
+    raise_warning('Failed to import tensorflow-probability: '
+                  'compilation of some probability distributions will fail.', 'red')
+    traceback.print_exc()
+    tfp = None
+
+
+def enumerate_literals(shape: Tuple[int, ...], axis: int, dtype: type=jnp.int32) -> jnp.ndarray:
     literals = jnp.arange(shape[axis], dtype=dtype)
     literals = literals[(...,) + (jnp.newaxis,) * (len(shape) - 1)]
     literals = jnp.moveaxis(literals, source=0, destination=axis)
@@ -74,7 +86,7 @@ class Comparison:
 class SigmoidComparison(Comparison):
     '''Comparison operations approximated using sigmoid functions.'''
     
-    def __init__(self, weight: float=10.0):
+    def __init__(self, weight: float=10.0) -> None:
         self.weight = weight
         
     # https://arxiv.org/abs/2110.05651
@@ -140,7 +152,7 @@ class Rounding:
 class SoftRounding(Rounding):
     '''Rounding operations approximated using soft operations.'''
     
-    def __init__(self, weight: float=10.0):
+    def __init__(self, weight: float=10.0) -> None:
         self.weight = weight
         
     # https://www.tensorflow.org/probability/api_docs/python/tfp/substrates/jax/bijectors/Softfloor
@@ -291,7 +303,7 @@ class YagerTNorm(TNorm):
     '''Yager t-norm given by the expression 
     (x, y) -> max(1 - ((1 - x)^p + (1 - y)^p)^(1/p)).'''
     
-    def __init__(self, p=2.0):
+    def __init__(self, p: float=2.0) -> None:
         self.p = float(p)
     
     def norm(self, id, init_params):
@@ -339,6 +351,9 @@ class RandomSampling:
     def binomial(self, id, init_params, logic):
         raise NotImplementedError
     
+    def negative_binomial(self, id, init_params, logic):
+        raise NotImplementedError
+    
     def geometric(self, id, init_params, logic):
         raise NotImplementedError
     
@@ -386,8 +401,7 @@ class SoftRandomSampling(RandomSampling):
     def _poisson_gumbel_softmax(self, id, init_params, logic):        
         argmax_approx = logic.argmax(id, init_params)
         def _jax_wrapped_calc_poisson_gumbel_softmax(key, rate, params):
-            ks = jnp.arange(0, self.poisson_bins)
-            ks = ks[(jnp.newaxis,) * jnp.ndim(rate) + (...,)]
+            ks = jnp.arange(self.poisson_bins)[(jnp.newaxis,) * jnp.ndim(rate) + (...,)]
             rate = rate[..., jnp.newaxis]
             log_prob = ks * jnp.log(rate + logic.eps) - rate - scipy.special.gammaln(ks + 1)
             Gumbel01 = random.gumbel(key=key, shape=jnp.shape(log_prob), dtype=logic.REAL)
@@ -400,10 +414,7 @@ class SoftRandomSampling(RandomSampling):
         less_approx = logic.less(id, init_params)
         def _jax_wrapped_calc_poisson_exponential(key, rate, params):
             Exp1 = random.exponential(
-                key=key, 
-                shape=(self.poisson_bins,) + jnp.shape(rate), 
-                dtype=logic.REAL
-            )
+                key=key,  shape=(self.poisson_bins,) + jnp.shape(rate), dtype=logic.REAL)
             delta_t = Exp1 / rate[jnp.newaxis, ...]
             times = jnp.cumsum(delta_t, axis=0)
             indicator, params = less_approx(times, 1.0, params)
@@ -411,72 +422,98 @@ class SoftRandomSampling(RandomSampling):
             return sample, params
         return _jax_wrapped_calc_poisson_exponential
 
+    # normal approximation to Poisson: Poisson(rate) -> Normal(rate, rate)
+    def _poisson_normal_approx(self, logic):
+        def _jax_wrapped_calc_poisson_normal_approx(key, rate, params):
+            normal = random.normal(key=key, shape=jnp.shape(rate), dtype=logic.REAL)
+            sample = rate + jnp.sqrt(rate) * normal
+            return sample, params    
+        return _jax_wrapped_calc_poisson_normal_approx
+    
     def poisson(self, id, init_params, logic):
-        def _jax_wrapped_calc_poisson_exact(key, rate, params):
-            sample = random.poisson(key=key, lam=rate, dtype=logic.INT)
-            sample = jnp.asarray(sample, dtype=logic.REAL)
-            return sample, params      
-              
         if self.poisson_exp_method:
             _jax_wrapped_calc_poisson_diff = self._poisson_exponential(
                 id, init_params, logic)
         else:
             _jax_wrapped_calc_poisson_diff = self._poisson_gumbel_softmax(
                 id, init_params, logic)
+        _jax_wrapped_calc_poisson_normal = self._poisson_normal_approx(logic)
         
+        # for small rate use the Poisson process or gumbel-softmax reparameterization
+        # for large rate use the normal approximation
         def _jax_wrapped_calc_poisson_approx(key, rate, params):
-            
-            # determine if error of truncation at rate is acceptable
             if self.poisson_bins > 0:
                 cuml_prob = scipy.stats.poisson.cdf(self.poisson_bins, rate)
-                approx_cond = jax.lax.stop_gradient(
-                    jnp.min(cuml_prob) > self.poisson_min_cdf)
+                small_rate = jax.lax.stop_gradient(cuml_prob >= self.poisson_min_cdf)
+                small_sample, params = _jax_wrapped_calc_poisson_diff(key, rate, params)
+                large_sample, params = _jax_wrapped_calc_poisson_normal(key, rate, params)
+                sample = jnp.where(small_rate, small_sample, large_sample)
+                return sample, params
             else:
-                approx_cond = False
-            
-            # for acceptable truncation use the approximation, use exact otherwise
-            return jax.lax.cond(
-                approx_cond, 
-                _jax_wrapped_calc_poisson_diff, 
-                _jax_wrapped_calc_poisson_exact, 
-                key, rate, params
-            )
+                return _jax_wrapped_calc_poisson_normal(key, rate, params)
         return _jax_wrapped_calc_poisson_approx        
     
-    def binomial(self, id, init_params, logic):
-        def _jax_wrapped_calc_binomial_exact(key, trials, prob, params):            
-            trials = jnp.asarray(trials, dtype=logic.REAL)
-            prob = jnp.asarray(prob, dtype=logic.REAL)
-            sample = random.binomial(key=key, n=trials, p=prob, dtype=logic.REAL)
-            return sample, params     
+    # normal approximation to Binomial: Bin(n, p) -> Normal(np, np(1-p))
+    def _binomial_normal_approx(self, logic):
+        def _jax_wrapped_calc_binomial_normal_approx(key, trials, prob, params):
+            normal = random.normal(key=key, shape=jnp.shape(trials), dtype=logic.REAL)
+            mean = trials * prob
+            std = jnp.sqrt(trials * prob * (1.0 - prob))
+            sample = mean + std * normal
+            return sample, params    
+        return _jax_wrapped_calc_binomial_normal_approx
         
-        # Binomial(n, p) = sum_{i = 1 ... n} Bernoulli(p)
-        bernoulli_approx = self.bernoulli(id, init_params, logic)
-        def _jax_wrapped_calc_binomial_sum(key, trials, prob, params):
-            prob_full = jnp.broadcast_to(
-                prob[..., jnp.newaxis], shape=jnp.shape(prob) + (self.binomial_bins,))
-            sample_bern, params = bernoulli_approx(key, prob_full, params)
-            indices = jnp.arange(self.binomial_bins)[
-                (jnp.newaxis,) * jnp.ndim(prob) + (...,)]
-            mask = indices < trials[..., jnp.newaxis]
-            sample = jnp.sum(sample_bern * mask, axis=-1)
-            return sample, params
+    def _binomial_gumbel_softmax(self, id, init_params, logic):
+        argmax_approx = logic.argmax(id, init_params)
+        def _jax_wrapped_calc_binomial_gumbel_softmax(key, trials, prob, params):
+            ks = jnp.arange(self.binomial_bins)[(jnp.newaxis,) * jnp.ndim(trials) + (...,)]
+            trials = trials[..., jnp.newaxis]
+            prob = prob[..., jnp.newaxis]
+            in_support = ks <= trials
+            ks = jnp.minimum(ks, trials)
+            log_prob = ((scipy.special.gammaln(trials + 1) - 
+                         scipy.special.gammaln(ks + 1) - 
+                         scipy.special.gammaln(trials - ks + 1)) +
+                        ks * jnp.log(prob + logic.eps) + 
+                        (trials - ks) * jnp.log1p(-prob + logic.eps))
+            log_prob = jnp.where(in_support, log_prob, jnp.log(logic.eps))
+            Gumbel01 = random.gumbel(key=key, shape=jnp.shape(log_prob), dtype=logic.REAL)
+            sample = Gumbel01 + log_prob
+            return argmax_approx(sample, axis=-1, params=params)
+        return _jax_wrapped_calc_binomial_gumbel_softmax
         
-        # for trials not too large use the Bernoulli relaxation, use exact otherwise
+    def binomial(self, id, init_params, logic):
+        _jax_wrapped_calc_binomial_normal = self._binomial_normal_approx(logic)  
+        _jax_wrapped_calc_binomial_gs = self._binomial_gumbel_softmax(id, init_params, logic)
+        
+        # for small trials use the Bernoulli relaxation
+        # for large trials use the normal approximation
         def _jax_wrapped_calc_binomial_approx(key, trials, prob, params):
-            return jax.lax.cond(
-                jax.lax.stop_gradient(jnp.max(trials) < self.binomial_bins), 
-                _jax_wrapped_calc_binomial_sum, 
-                _jax_wrapped_calc_binomial_exact, 
-                key, trials, prob, params
-            )
+            small_trials = jax.lax.stop_gradient(trials < self.binomial_bins)
+            small_sample, params = _jax_wrapped_calc_binomial_gs(key, trials, prob, params)
+            large_sample, params = _jax_wrapped_calc_binomial_normal(key, trials, prob, params)
+            sample = jnp.where(small_trials, small_sample, large_sample)
+            return sample, params
         return _jax_wrapped_calc_binomial_approx
     
+    # https://en.wikipedia.org/wiki/Negative_binomial_distribution#Gamma%E2%80%93Poisson_mixture
+    def negative_binomial(self, id, init_params, logic):
+        poisson_approx = self.poisson(id, init_params, logic)
+        def _jax_wrapped_calc_negative_binomial_approx(key, trials, prob, params):
+            key, subkey = random.split(key)
+            trials = jnp.asarray(trials, dtype=logic.REAL)
+            Gamma = random.gamma(key=key, a=trials, dtype=logic.REAL)
+            scale = (1.0 - prob) / prob
+            poisson_rate = scale * Gamma
+            return poisson_approx(subkey, poisson_rate, params)
+        return _jax_wrapped_calc_negative_binomial_approx
+
     def geometric(self, id, init_params, logic):
         approx_floor = logic.floor(id, init_params)
         def _jax_wrapped_calc_geometric_approx(key, prob, params):
             U = random.uniform(key=key, shape=jnp.shape(prob), dtype=logic.REAL)
-            floor, params = approx_floor(jnp.log1p(-U) / jnp.log1p(-prob), params)
+            floor, params = approx_floor(
+                jnp.log1p(-U) / jnp.log1p(-prob + logic.eps), params)
             sample = floor + 1
             return sample, params
         return _jax_wrapped_calc_geometric_approx
@@ -532,6 +569,14 @@ class Determinization(RandomSampling):
     def binomial(self, id, init_params, logic):
         return self._jax_wrapped_calc_binomial_determinized
     
+    @staticmethod
+    def _jax_wrapped_calc_negative_binomial_determinized(key, trials, prob, params):
+        sample = trials * ((1.0 / prob) - 1.0)
+        return sample, params
+
+    def negative_binomial(self, id, init_params, logic):
+        return self._jax_wrapped_calc_negative_binomial_determinized
+    
     @staticmethod
     def _jax_wrapped_calc_geometric_determinized(key, prob, params):
         sample = 1.0 / prob
@@ -712,7 +757,8 @@ class Logic:
                 'Discrete': self.discrete,
                 'Poisson': self.poisson,
                 'Geometric': self.geometric,
-                'Binomial': self.binomial
+                'Binomial': self.binomial,
+                'NegativeBinomial': self.negative_binomial
             }
         }
 
@@ -830,6 +876,9 @@ class Logic:
     def binomial(self, id, init_params):
         raise NotImplementedError
 
+    def negative_binomial(self, id, init_params):
+        raise NotImplementedError
+
 
 class ExactLogic(Logic):
     '''A class representing exact logic in JAX.'''
@@ -1005,6 +1054,17 @@ class ExactLogic(Logic):
             sample = jnp.asarray(sample, dtype=self.INT)
             return sample, params
         return _jax_wrapped_calc_binomial_exact
+    
+    # note: for some reason tfp defines it as number of successes before trials failures
+    # I will define it as the number of failures before trials successes
+    def negative_binomial(self, id, init_params):
+        def _jax_wrapped_calc_negative_binomial_exact(key, trials, prob, params):
+            trials = jnp.asarray(trials, dtype=self.REAL)
+            prob = jnp.asarray(prob, dtype=self.REAL)
+            dist = tfp.distributions.NegativeBinomial(total_count=trials, probs=1.0 - prob)
+            sample = jnp.asarray(dist.sample(seed=key), dtype=self.INT)
+            return sample, params
+        return _jax_wrapped_calc_negative_binomial_exact
 
     
 class FuzzyLogic(Logic):
@@ -1234,6 +1294,9 @@ class FuzzyLogic(Logic):
     
     def binomial(self, id, init_params):
         return self.sampling.binomial(id, init_params, self)
+    
+    def negative_binomial(self, id, init_params):
+        return self.sampling.negative_binomial(id, init_params, self)
 
 
 # ===========================================================================

{pyrddlgym_jax-2.2 → pyrddlgym_jax-2.3}/pyRDDLGym_jax.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: pyRDDLGym-jax
-Version: 2.2
+Version: 2.3
 Summary: pyRDDLGym-jax: automatic differentiation for solving sequential planning problems in JAX.
 Home-page: https://github.com/pyrddlgym-project/pyRDDLGym-jax
 Author: Michael Gimelfarb, Ayal Taitler, Scott Sanner

{pyrddlgym_jax-2.2 → pyrddlgym_jax-2.3}/setup.py

@@ -19,7 +19,7 @@ long_description = (Path(__file__).parent / "README.md").read_text()
 
 setup(
       name='pyRDDLGym-jax',
-      version='2.2',
+      version='2.3',
       author="Michael Gimelfarb, Ayal Taitler, Scott Sanner",
       author_email="mike.gimelfarb@mail.utoronto.ca, ataitler@gmail.com, ssanner@mie.utoronto.ca",
       description="pyRDDLGym-jax: automatic differentiation for solving sequential planning problems in JAX.",

pyrddlgym_jax-2.2/pyRDDLGym_jax/__init__.py

@@ -1 +0,0 @@
-__version__ = '2.2'