pyRDDLGym-jax 2.0__py3-none-any.whl → 2.1__py3-none-any.whl

pyRDDLGym_jax/core/logic.py

@@ -24,12 +24,17 @@
 # Gumble-Softmax." In International Conference on Learning Representations (ICLR 2017). 
 # OpenReview. net, 2017.
 #
+# [6] Vafaii, H., Galor, D., & Yates, J. (2025). Poisson Variational Autoencoder. 
+# Advances in Neural Information Processing Systems, 37, 44871-44906.
+#
 # ***********************************************************************
 
+from typing import Callable, Dict, Union
 
 import jax
 import jax.numpy as jnp
 import jax.random as random
+import jax.scipy as scipy 
 
 
 def enumerate_literals(shape, axis, dtype=jnp.int32):
@@ -105,7 +110,7 @@ class SigmoidComparison(Comparison):
         id_ = str(id)
         init_params[id_] = self.weight
         def _jax_wrapped_calc_argmax_approx(x, axis, params):
-            literals = enumerate_literals(x.shape, axis=axis)
+            literals = enumerate_literals(jnp.shape(x), axis=axis)
             softmax = jax.nn.softmax(params[id_] * x, axis=axis)
             sample = jnp.sum(literals * softmax, axis=axis)
             return sample, params
@@ -323,62 +328,192 @@ class YagerTNorm(TNorm):
 # ===========================================================================
 
 class RandomSampling:
-    '''An abstract class that describes how discrete and non-reparameterizable 
-    random variables are sampled.'''
+    '''Describes how non-reparameterizable random variables are sampled.'''
     
     def discrete(self, id, init_params, logic):
         raise NotImplementedError
     
-    def bernoulli(self, id, init_params, logic):
-        discrete_approx = self.discrete(id, init_params, logic)        
-        def _jax_wrapped_calc_bernoulli_approx(key, prob, params):
-            prob = jnp.stack([1.0 - prob, prob], axis=-1)
-            return discrete_approx(key, prob, params)        
-        return _jax_wrapped_calc_bernoulli_approx
-    
-    @staticmethod
-    def _jax_wrapped_calc_poisson_exact(key, rate, params):
-        sample = random.poisson(key=key, lam=rate, dtype=logic.INT)
-        return sample, params
-        
     def poisson(self, id, init_params, logic):
-        return self._jax_wrapped_calc_poisson_exact
+        raise NotImplementedError
+    
+    def binomial(self, id, init_params, logic):
+        raise NotImplementedError
     
     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.log(U) / jnp.log(1.0 - prob), params)
-            sample = floor + 1
-            return sample, params
-        return _jax_wrapped_calc_geometric_approx
+        raise NotImplementedError
+    
+    def bernoulli(self, id, init_params, logic):
+        raise NotImplementedError
+    
+    def __str__(self) -> str:
+        return 'RandomSampling'
 
 
-class GumbelSoftmax(RandomSampling):
+class SoftRandomSampling(RandomSampling):
     '''Random sampling of discrete variables using Gumbel-softmax trick.'''
     
+    def __init__(self, poisson_max_bins: int=100, 
+                 poisson_min_cdf: float=0.999,
+                 poisson_exp_sampling: bool=True,
+                 binomial_max_bins: int=100,
+                 bernoulli_gumbel_softmax: bool=False) -> None:
+        '''Creates a new instance of soft random sampling.
+
+        :param poisson_max_bins: maximum bins to use for Poisson distribution relaxation
+        :param poisson_min_cdf: minimum cdf value of Poisson within truncated region
+        in order to use Poisson relaxation
+        :param poisson_exp_sampling: whether to use Poisson process sampling method
+        instead of truncated Gumbel-Softmax
+        :param binomial_max_bins: maximum bins to use for Binomial distribution relaxation
+        :param bernoulli_gumbel_softmax: whether to use Gumbel-Softmax to approximate
+        Bernoulli samples, or the standard uniform reparameterization instead
+        '''
+        self.poisson_bins = poisson_max_bins
+        self.poisson_min_cdf = poisson_min_cdf
+        self.poisson_exp_method = poisson_exp_sampling
+        self.binomial_bins = binomial_max_bins
+        self.bernoulli_gumbel_softmax = bernoulli_gumbel_softmax
+
     # https://arxiv.org/pdf/1611.01144
     def discrete(self, id, init_params, logic):
         argmax_approx = logic.argmax(id, init_params)
         def _jax_wrapped_calc_discrete_gumbel_softmax(key, prob, params):
-            Gumbel01 = random.gumbel(key=key, shape=prob.shape, dtype=logic.REAL)
+            Gumbel01 = random.gumbel(key=key, shape=jnp.shape(prob), dtype=logic.REAL)
             sample = Gumbel01 + jnp.log(prob + logic.eps)
             return argmax_approx(sample, axis=-1, params=params)
         return _jax_wrapped_calc_discrete_gumbel_softmax
     
+    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) + (...,)]
+            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)
+            sample = Gumbel01 + log_prob
+            return argmax_approx(sample, axis=-1, params=params)
+        return _jax_wrapped_calc_poisson_gumbel_softmax
+    
+    # https://arxiv.org/abs/2405.14473
+    def _poisson_exponential(self, id, init_params, logic):
+        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
+            )
+            delta_t = Exp1 / rate[jnp.newaxis, ...]
+            times = jnp.cumsum(delta_t, axis=0)
+            indicator, params = less_approx(times, 1.0, params)
+            sample = jnp.sum(indicator, axis=0)
+            return sample, params
+        return _jax_wrapped_calc_poisson_exponential
+
+    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)
+        
+        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)
+            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_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     
+        
+        # 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
+        
+        # for trials not too large use the Bernoulli relaxation, use exact otherwise
+        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
+            )
+        return _jax_wrapped_calc_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)
+            sample = floor + 1
+            return sample, params
+        return _jax_wrapped_calc_geometric_approx
+    
+    def _bernoulli_uniform(self, id, init_params, logic):
+        less_approx = logic.less(id, init_params)  
+        def _jax_wrapped_calc_bernoulli_uniform(key, prob, params):
+            U = random.uniform(key=key, shape=jnp.shape(prob), dtype=logic.REAL)
+            return less_approx(U, prob, params)   
+        return _jax_wrapped_calc_bernoulli_uniform
+    
+    def _bernoulli_gumbel_softmax(self, id, init_params, logic):
+        discrete_approx = self.discrete(id, init_params, logic)        
+        def _jax_wrapped_calc_bernoulli_gumbel_softmax(key, prob, params):
+            prob = jnp.stack([1.0 - prob, prob], axis=-1)
+            return discrete_approx(key, prob, params)        
+        return _jax_wrapped_calc_bernoulli_gumbel_softmax
+    
+    def bernoulli(self, id, init_params, logic):
+        if self.bernoulli_gumbel_softmax:
+            return self._bernoulli_gumbel_softmax(id, init_params, logic)
+        else:
+            return self._bernoulli_uniform(id, init_params, logic)
+
     def __str__(self) -> str:
-        return 'Gumbel-Softmax'
+        return 'SoftRandomSampling'
     
 
 class Determinization(RandomSampling):
     '''Random sampling of variables using their deterministic mean estimate.'''
-    
+
     @staticmethod
     def _jax_wrapped_calc_discrete_determinized(key, prob, params):
-        literals = enumerate_literals(prob.shape, axis=-1)
+        literals = enumerate_literals(jnp.shape(prob), axis=-1)
         sample = jnp.sum(literals * prob, axis=-1)
         return sample, params
-        
+    
     def discrete(self, id, init_params, logic):
         return self._jax_wrapped_calc_discrete_determinized
     
@@ -389,6 +524,14 @@ class Determinization(RandomSampling):
     def poisson(self, id, init_params, logic):
         return self._jax_wrapped_calc_poisson_determinized
     
+    @staticmethod
+    def _jax_wrapped_calc_binomial_determinized(key, trials, prob, params):
+        sample = trials * prob
+        return sample, params
+    
+    def binomial(self, id, init_params, logic):
+        return self._jax_wrapped_calc_binomial_determinized
+    
     @staticmethod
     def _jax_wrapped_calc_geometric_determinized(key, prob, params):
         sample = 1.0 / prob
@@ -397,6 +540,14 @@ class Determinization(RandomSampling):
     def geometric(self, id, init_params, logic):
         return self._jax_wrapped_calc_geometric_determinized
     
+    @staticmethod
+    def _jax_wrapped_calc_bernoulli_determinized(key, prob, params):
+        sample = prob
+        return sample, params
+    
+    def bernoulli(self, id, init_params, logic):
+        return self._jax_wrapped_calc_bernoulli_determinized
+
     def __str__(self) -> str:
         return 'Deterministic'
     
@@ -435,8 +586,8 @@ class SoftControlFlow(ControlFlow):
         id_ = str(id)
         init_params[id_] = self.weight
         def _jax_wrapped_calc_switch_soft(pred, cases, params):
-            literals = enumerate_literals(cases.shape, axis=0)
-            pred = jnp.broadcast_to(pred[jnp.newaxis, ...], shape=cases.shape)
+            literals = enumerate_literals(jnp.shape(cases), axis=0)
+            pred = jnp.broadcast_to(pred[jnp.newaxis, ...], shape=jnp.shape(cases))
             proximity = -jnp.square(pred - literals)
             softcase = jax.nn.softmax(params[id_] * proximity, axis=0)
             sample = jnp.sum(cases * softcase, axis=0)
@@ -477,6 +628,94 @@ class Logic:
             self.INT = jnp.int32
             jax.config.update('jax_enable_x64', False)
     
+    @staticmethod
+    def wrap_logic(func):
+        def exact_func(id, init_params):
+            return func
+        return exact_func
+        
+    def get_operator_dicts(self) -> Dict[str, Union[Callable, Dict[str, Callable]]]:
+        '''Returns a dictionary of all operators in the current logic.'''
+        return {
+            'negative': self.wrap_logic(ExactLogic.exact_unary_function(jnp.negative)),
+            'arithmetic': {
+                '+': self.wrap_logic(ExactLogic.exact_binary_function(jnp.add)),
+                '-': self.wrap_logic(ExactLogic.exact_binary_function(jnp.subtract)),
+                '*': self.wrap_logic(ExactLogic.exact_binary_function(jnp.multiply)),
+                '/': self.wrap_logic(ExactLogic.exact_binary_function(jnp.divide))
+            },
+            'relational': {
+                '>=': self.greater_equal,
+                '<=': self.less_equal,
+                '<': self.less,
+                '>': self.greater,
+                '==': self.equal,
+                '~=': self.not_equal
+            },
+            'logical_not': self.logical_not,
+            'logical': {
+                '^': self.logical_and,
+                '&': self.logical_and,
+                '|': self.logical_or,
+                '~': self.xor,
+                '=>': self.implies,
+                '<=>': self.equiv
+            },
+            'aggregation': {
+                'sum': self.wrap_logic(ExactLogic.exact_aggregation(jnp.sum)),
+                'avg': self.wrap_logic(ExactLogic.exact_aggregation(jnp.mean)),
+                'prod': self.wrap_logic(ExactLogic.exact_aggregation(jnp.prod)),
+                'minimum': self.wrap_logic(ExactLogic.exact_aggregation(jnp.min)),
+                'maximum': self.wrap_logic(ExactLogic.exact_aggregation(jnp.max)),
+                'forall': self.forall,
+                'exists': self.exists,
+                'argmin': self.argmin,
+                'argmax': self.argmax
+            },
+            'unary': {
+                'abs': self.wrap_logic(ExactLogic.exact_unary_function(jnp.abs)),
+                'sgn': self.sgn,
+                'round': self.round,
+                'floor': self.floor,
+                'ceil': self.ceil,
+                'cos': self.wrap_logic(ExactLogic.exact_unary_function(jnp.cos)),
+                'sin': self.wrap_logic(ExactLogic.exact_unary_function(jnp.sin)),
+                'tan': self.wrap_logic(ExactLogic.exact_unary_function(jnp.tan)),
+                'acos': self.wrap_logic(ExactLogic.exact_unary_function(jnp.arccos)),
+                'asin': self.wrap_logic(ExactLogic.exact_unary_function(jnp.arcsin)),
+                'atan': self.wrap_logic(ExactLogic.exact_unary_function(jnp.arctan)),
+                'cosh': self.wrap_logic(ExactLogic.exact_unary_function(jnp.cosh)),
+                'sinh': self.wrap_logic(ExactLogic.exact_unary_function(jnp.sinh)),
+                'tanh': self.wrap_logic(ExactLogic.exact_unary_function(jnp.tanh)),
+                'exp': self.wrap_logic(ExactLogic.exact_unary_function(jnp.exp)),
+                'ln': self.wrap_logic(ExactLogic.exact_unary_function(jnp.log)),
+                'sqrt': self.sqrt,
+                'lngamma': self.wrap_logic(ExactLogic.exact_unary_function(scipy.special.gammaln)),
+                'gamma': self.wrap_logic(ExactLogic.exact_unary_function(scipy.special.gamma))
+            },
+            'binary': {
+                'div': self.div,
+                'mod': self.mod,
+                'fmod': self.mod,
+                'min': self.wrap_logic(ExactLogic.exact_binary_function(jnp.minimum)),
+                'max': self.wrap_logic(ExactLogic.exact_binary_function(jnp.maximum)),
+                'pow': self.wrap_logic(ExactLogic.exact_binary_function(jnp.power)),
+                'log': self.wrap_logic(ExactLogic.exact_binary_log),
+                'hypot': self.wrap_logic(ExactLogic.exact_binary_function(jnp.hypot)),
+            },
+            'control': {
+                'if': self.control_if,
+                'switch': self.control_switch
+            },
+            'sampling': {
+                'Bernoulli': self.bernoulli,
+                'Discrete': self.discrete,
+                'Poisson': self.poisson,
+                'Geometric': self.geometric,
+                'Binomial': self.binomial
+            }
+        }
+
     # ===========================================================================
     # logical operators
     # ===========================================================================
@@ -587,6 +826,9 @@ class Logic:
     
     def geometric(self, id, init_params):
         raise NotImplementedError
+    
+    def binomial(self, id, init_params):
+        raise NotImplementedError
 
 
 class ExactLogic(Logic):
@@ -627,11 +869,11 @@ class ExactLogic(Logic):
         return self.exact_binary_function(jnp.logical_xor)
     
     @staticmethod
-    def exact_binary_implies(x, y, params):
+    def _jax_wrapped_calc_implies_exact(x, y, params):
         return jnp.logical_or(jnp.logical_not(x), y), params     
     
     def implies(self, id, init_params):
-        return self.exact_binary_implies
+        return self._jax_wrapped_calc_implies_exact
     
     def equiv(self, id, init_params):
         return self.exact_binary_function(jnp.equal)
@@ -668,6 +910,10 @@ class ExactLogic(Logic):
     # special functions
     # ===========================================================================
      
+    @staticmethod
+    def exact_binary_log(x, y, params):
+        return jnp.log(x) / jnp.log(y), params
+    
     def sgn(self, id, init_params):
         return self.exact_unary_function(jnp.sign)
     
@@ -704,54 +950,62 @@ class ExactLogic(Logic):
     # ===========================================================================
     
     @staticmethod
-    def exact_if_then_else(c, a, b, params):
+    def _jax_wrapped_calc_if_then_else_exact(c, a, b, params):
         return jnp.where(c > 0.5, a, b), params
          
     def control_if(self, id, init_params):
-        return self.exact_if_then_else
+        return self._jax_wrapped_calc_if_then_else_exact
     
     @staticmethod
-    def exact_switch(pred, cases, params):
+    def _jax_wrapped_calc_switch_exact(pred, cases, params):
         pred = pred[jnp.newaxis, ...]
         sample = jnp.take_along_axis(cases, pred, axis=0)
         assert sample.shape[0] == 1
         return sample[0, ...], params
     
     def control_switch(self, id, init_params):
-        return self.exact_switch
+        return self._jax_wrapped_calc_switch_exact
     
     # ===========================================================================
     # random variables
     # ===========================================================================
     
     @staticmethod
-    def exact_discrete(key, prob, params):
-        return random.categorical(key=key, logits=jnp.log(prob), axis=-1), params  
+    def _jax_wrapped_calc_discrete_exact(key, prob, params):
+        sample = random.categorical(key=key, logits=jnp.log(prob), axis=-1)
+        return sample, params  
     
     def discrete(self, id, init_params):
-        return self.exact_discrete
+        return self._jax_wrapped_calc_discrete_exact
     
     @staticmethod
-    def exact_bernoulli(key, prob, params):
+    def _jax_wrapped_calc_bernoulli_exact(key, prob, params):
         return random.bernoulli(key, prob), params
     
     def bernoulli(self, id, init_params):
-        return self.exact_bernoulli
-    
-    @staticmethod
-    def exact_poisson(key, rate, params):
-        return random.poisson(key=key, lam=rate), params   
+        return self._jax_wrapped_calc_bernoulli_exact
     
     def poisson(self, id, init_params):
-        return self.exact_poisson
-    
-    @staticmethod
-    def exact_geometric(key, prob, params):
-        return random.geometric(key=key, p=prob), params
+        def _jax_wrapped_calc_poisson_exact(key, rate, params):
+            sample = random.poisson(key=key, lam=rate, dtype=self.INT)
+            return sample, params
+        return _jax_wrapped_calc_poisson_exact
     
     def geometric(self, id, init_params):
-        return self.exact_geometric
-    
+        def _jax_wrapped_calc_geometric_exact(key, prob, params):
+            sample = random.geometric(key=key, p=prob, dtype=self.INT)
+            return sample, params
+        return _jax_wrapped_calc_geometric_exact
+    
+    def binomial(self, id, init_params):
+        def _jax_wrapped_calc_binomial_exact(key, trials, prob, params):
+            trials = jnp.asarray(trials, dtype=self.REAL)
+            prob = jnp.asarray(prob, dtype=self.REAL)
+            sample = random.binomial(key=key, n=trials, p=prob, dtype=self.REAL)
+            sample = jnp.asarray(sample, dtype=self.INT)
+            return sample, params
+        return _jax_wrapped_calc_binomial_exact
+
     
 class FuzzyLogic(Logic):
     '''A class representing fuzzy logic in JAX.'''
@@ -759,7 +1013,7 @@ class FuzzyLogic(Logic):
     def __init__(self, tnorm: TNorm=ProductTNorm(),
                  complement: Complement=StandardComplement(),
                  comparison: Comparison=SigmoidComparison(),
-                 sampling: RandomSampling=GumbelSoftmax(),
+                 sampling: RandomSampling=SoftRandomSampling(),
                  rounding: Rounding=SoftRounding(),
                  control: ControlFlow=SoftControlFlow(),
                  eps: float=1e-15,
@@ -977,6 +1231,9 @@ class FuzzyLogic(Logic):
     
     def geometric(self, id, init_params):
         return self.sampling.geometric(id, init_params, self)
+    
+    def binomial(self, id, init_params):
+        return self.sampling.binomial(id, init_params, self)
 
 
 # ===========================================================================
@@ -1011,8 +1268,8 @@ def _test_logical():
         pred, w = _if(cond, +1, -1, w)
         return pred
     
-    x1 = jnp.asarray([1, 1, -1, -1, 0.1, 15, -0.5]).astype(float)
-    x2 = jnp.asarray([1, -1, 1, -1, 10, -30, 6]).astype(float)
+    x1 = jnp.asarray([1, 1, -1, -1, 0.1, 15, -0.5], dtype=float)
+    x2 = jnp.asarray([1, -1, 1, -1, 10, -30, 6], dtype=float)
     print(test_logic(x1, x2, init_params))