pyRDDLGym-jax 0.1__py3-none-any.whl → 0.2__py3-none-any.whl

pyRDDLGym_jax/core/logic.py

@@ -1,7 +1,7 @@
 import jax
 import jax.numpy as jnp
 import jax.random as random
-from typing import Set
+from typing import Optional, Set
 
 from pyRDDLGym.core.debug.exception import raise_warning
 
@@ -20,6 +20,32 @@ class StandardComplement(Complement):
         return 1.0 - x
 
 
+class Comparison:
+    '''Base class for approximate comparison operations.'''
+    
+    def greater_equal(self, x, y, param):
+        raise NotImplementedError
+    
+    def greater(self, x, y, param):
+        raise NotImplementedError
+    
+    def equal(self, x, y, param):
+        raise NotImplementedError
+    
+
+class SigmoidComparison(Comparison):
+    '''Comparison operations approximated using sigmoid functions.'''
+    
+    def greater_equal(self, x, y, param):
+        return jax.nn.sigmoid(param * (x - y))
+    
+    def greater(self, x, y, param):
+        return jax.nn.sigmoid(param * (x - y))
+    
+    def equal(self, x, y, param):
+        return 1.0 - jnp.square(jnp.tanh(param * (y - x)))
+        
+
 class TNorm:
     '''Base class for fuzzy differentiable t-norms.'''
     
@@ -71,40 +97,61 @@ class FuzzyLogic:
     
     def __init__(self, tnorm: TNorm=ProductTNorm(),
                  complement: Complement=StandardComplement(),
+                 comparison: Comparison=SigmoidComparison(),
                  weight: float=10.0,
-                 debias: Set[str]={},
+                 debias: Optional[Set[str]]=None,
                  eps: float=1e-10,
-                 verbose: bool=False):
+                 verbose: bool=False,
+                 use64bit: bool=False) -> None:
         '''Creates a new fuzzy logic in Jax.
         
         :param tnorm: fuzzy operator for logical AND
         :param complement: fuzzy operator for logical NOT
-        :param weight: a concentration parameter (larger means better accuracy)
+        :param comparison: fuzzy operator for comparisons (>, >=, <, ==, ~=, ...)
+        :param weight: a sharpness parameter for sigmoid and softmax activations
         :param error: an error parameter (e.g. floor) (smaller means better accuracy)
         :param debias: which functions to de-bias approximate on forward pass
         :param eps: small positive float to mitigate underflow
         :param verbose: whether to dump replacements and other info to console
+        :param use64bit: whether to perform arithmetic in 64 bit
         '''
         self.tnorm = tnorm
         self.complement = complement
+        self.comparison = comparison
         self.weight = float(weight)
+        if debias is None:
+            debias = set()
         self.debias = debias
         self.eps = eps
         self.verbose = verbose
-    
-    def summarize_hyperparameters(self):
+        self.set_use64bit(use64bit)
+    
+    def set_use64bit(self, use64bit: bool) -> None:
+        self.use64bit = use64bit
+        if use64bit:
+            self.REAL = jnp.float64
+            self.INT = jnp.int64
+            jax.config.update('jax_enable_x64', True)
+        else:
+            self.REAL = jnp.float32
+            self.INT = jnp.int32
+            jax.config.update('jax_enable_x64', False)
+        
+    def summarize_hyperparameters(self) -> None:
         print(f'model relaxation:\n'
               f'    tnorm         ={type(self.tnorm).__name__}\n'
               f'    complement    ={type(self.complement).__name__}\n'
+              f'    comparison    ={type(self.comparison).__name__}\n'
               f'    sigmoid_weight={self.weight}\n'
               f'    cpfs_to_debias={self.debias}\n'
-              f'    underflow_tol ={self.eps}')
+              f'    underflow_tol ={self.eps}\n'
+              f'    use64bit      ={self.use64bit}')
         
     # ===========================================================================
     # logical operators
     # ===========================================================================
      
-    def And(self):
+    def logical_and(self):
         if self.verbose:
             raise_warning('Using the replacement rule: a ^ b --> tnorm(a, b).')
         
@@ -115,9 +162,9 @@ class FuzzyLogic:
         
         return _jax_wrapped_calc_and_approx, None
     
-    def Not(self):
+    def logical_not(self):
         if self.verbose:
-            raise_warning('Using the replacement rule: ~a --> 1 - a')
+            raise_warning('Using the replacement rule: ~a --> complement(a)')
         
         _not = self.complement
         
@@ -126,9 +173,9 @@ class FuzzyLogic:
         
         return _jax_wrapped_calc_not_approx, None
     
-    def Or(self):
+    def logical_or(self):
         if self.verbose:
-            raise_warning('Using the replacement rule: a or b --> tconorm(a, b).')
+            raise_warning('Using the replacement rule: a | b --> tconorm(a, b).')
             
         _not = self.complement
         _and = self.tnorm.norm
@@ -141,7 +188,7 @@ class FuzzyLogic:
     def xor(self):
         if self.verbose:
             raise_warning('Using the replacement rule: '
-                          'a xor b --> (a or b) ^ (a ^ b).')
+                          'a ~ b --> (a | b) ^ (a ^ b).')
         
         _not = self.complement
         _and = self.tnorm.norm
@@ -182,7 +229,7 @@ class FuzzyLogic:
     def forall(self):
         if self.verbose:
             raise_warning('Using the replacement rule: '
-                          'forall(a) --> tnorm(a[1], tnorm(a[2], ...))')
+                          'forall(a) --> a[1] ^ a[2] ^ ...')
         
         _forall = self.tnorm.norms
         
@@ -204,31 +251,35 @@ class FuzzyLogic:
     # comparison operators
     # ===========================================================================
      
-    def greaterEqual(self):
+    def greater_equal(self):
         if self.verbose:
-            raise_warning('Using the replacement rule: a >= b --> sigmoid(a - b)')
-            
-        debias = 'greaterEqual' in self.debias
+            raise_warning('Using the replacement rule: '
+                          'a >= b --> comparison.greater_equal(a, b)')
+        
+        greater_equal_op = self.comparison.greater_equal
+        debias = 'greater_equal' in self.debias
         
         def _jax_wrapped_calc_geq_approx(a, b, param):
-            sample = jax.nn.sigmoid(param * (a - b))
+            sample = greater_equal_op(a, b, param)
             if debias:
                 hard_sample = jnp.greater_equal(a, b)
                 sample += jax.lax.stop_gradient(hard_sample - sample)
             return sample
         
-        tags = ('weight', 'greaterEqual')
+        tags = ('weight', 'greater_equal')
         new_param = (tags, self.weight)
         return _jax_wrapped_calc_geq_approx, new_param
     
     def greater(self):
         if self.verbose:
-            raise_warning('Using the replacement rule: a > b --> sigmoid(a - b)')
-            
+            raise_warning('Using the replacement rule: '
+                          'a > b --> comparison.greater(a, b)')
+        
+        greater_op = self.comparison.greater
         debias = 'greater' in self.debias
         
         def _jax_wrapped_calc_gre_approx(a, b, param):
-            sample = jax.nn.sigmoid(param * (a - b))
+            sample = greater_op(a, b, param)
             if debias:
                 hard_sample = jnp.greater(a, b)
                 sample += jax.lax.stop_gradient(hard_sample - sample)
@@ -238,8 +289,8 @@ class FuzzyLogic:
         new_param = (tags, self.weight)
         return _jax_wrapped_calc_gre_approx, new_param
     
-    def lessEqual(self):
-        jax_geq, jax_param = self.greaterEqual()
+    def less_equal(self):
+        jax_geq, jax_param = self.greater_equal()
         
         def _jax_wrapped_calc_leq_approx(a, b, param):
             return jax_geq(-a, -b, param)
@@ -256,12 +307,14 @@ class FuzzyLogic:
 
     def equal(self):
         if self.verbose:
-            raise_warning('Using the replacement rule: a == b --> sech^2(b - a)')
-            
+            raise_warning('Using the replacement rule: '
+                          'a == b --> comparison.equal(a, b)')
+        
+        equal_op = self.comparison.equal
         debias = 'equal' in self.debias
         
         def _jax_wrapped_calc_equal_approx(a, b, param):
-            sample = 1.0 - jnp.square(jnp.tanh(param * (b - a)))
+            sample = equal_op(a, b, param)
             if debias:
                 hard_sample = jnp.equal(a, b)
                 sample += jax.lax.stop_gradient(hard_sample - sample)
@@ -271,7 +324,7 @@ class FuzzyLogic:
         new_param = (tags, self.weight)
         return _jax_wrapped_calc_equal_approx, new_param
     
-    def notEqual(self):
+    def not_equal(self):
         _not = self.complement
         jax_eq, jax_param = self.equal()
         
@@ -284,31 +337,32 @@ class FuzzyLogic:
     # special functions
     # ===========================================================================
      
-    def signum(self):
+    def sgn(self):
         if self.verbose:
-            raise_warning('Using the replacement rule: signum(x) --> tanh(x)')
+            raise_warning('Using the replacement rule: sgn(x) --> tanh(x)')
             
-        debias = 'signum' in self.debias
+        debias = 'sgn' in self.debias
         
-        def _jax_wrapped_calc_signum_approx(x, param):
+        def _jax_wrapped_calc_sgn_approx(x, param):
             sample = jnp.tanh(param * x)
             if debias:
                 hard_sample = jnp.sign(x)
                 sample += jax.lax.stop_gradient(hard_sample - sample)
             return sample
         
-        tags = ('weight', 'signum')
+        tags = ('weight', 'sgn')
         new_param = (tags, self.weight)
-        return _jax_wrapped_calc_signum_approx, new_param
+        return _jax_wrapped_calc_sgn_approx, new_param
     
     def floor(self):
         if self.verbose:
             raise_warning('Using the replacement rule: '
                           'floor(x) --> x - atan(-1.0 / tan(pi * x)) / pi - 0.5')
-            
+        
         def _jax_wrapped_calc_floor_approx(x, param):
             sawtooth_part = jnp.arctan(-1.0 / jnp.tan(x * jnp.pi)) / jnp.pi + 0.5
-            return x - jax.lax.stop_gradient(sawtooth_part)
+            sample = x - jax.lax.stop_gradient(sawtooth_part)
+            return sample
         
         return _jax_wrapped_calc_floor_approx, None
         
@@ -324,8 +378,14 @@ class FuzzyLogic:
         if self.verbose:
             raise_warning('Using the replacement rule: round(x) --> x')
         
+        debias = 'round' in self.debias
+        
         def _jax_wrapped_calc_round_approx(x, param):
-            return x
+            sample = x
+            if debias:
+                hard_sample = jnp.round(x)
+                sample += jax.lax.stop_gradient(hard_sample - sample)
+            return sample
         
         return _jax_wrapped_calc_round_approx, None
     
@@ -337,13 +397,13 @@ class FuzzyLogic:
         
         return _jax_wrapped_calc_mod_approx, jax_param
     
-    def floorDiv(self):
+    def div(self):
         jax_floor, jax_param = self.floor()
         
-        def _jax_wrapped_calc_mod_approx(x, y, param):
+        def _jax_wrapped_calc_div_approx(x, y, param):
             return jax_floor(x / y, param)
         
-        return _jax_wrapped_calc_mod_approx, jax_param
+        return _jax_wrapped_calc_div_approx, jax_param
     
     def sqrt(self):
         if self.verbose:
@@ -369,14 +429,14 @@ class FuzzyLogic:
     def argmax(self):
         if self.verbose:
             raise_warning('Using the replacement rule: '
-                          f'argmax(x) --> sum(i * softmax(x[i]))')
+                          'argmax(x) --> sum(i * softmax(x[i]))')
             
         debias = 'argmax' in self.debias
         
         def _jax_wrapped_calc_argmax_approx(x, axis, param):
-            prob_max = jax.nn.softmax(param * x, axis=axis)
-            literals = FuzzyLogic._literals(prob_max.shape, axis=axis)
-            sample = jnp.sum(literals * prob_max, axis=axis)
+            literals = FuzzyLogic._literals(x.shape, axis=axis)
+            soft_max = jax.nn.softmax(param * x, axis=axis)
+            sample = jnp.sum(literals * soft_max, axis=axis)
             if debias:
                 hard_sample = jnp.argmax(x, axis=axis)
                 sample += jax.lax.stop_gradient(hard_sample - sample)
@@ -398,26 +458,33 @@ class FuzzyLogic:
     # control flow
     # ===========================================================================
      
-    def If(self):
+    def control_if(self):
         if self.verbose:
             raise_warning('Using the replacement rule: '
                           'if c then a else b --> c * a + (1 - c) * b')
         
+        debias = 'if' in self.debias
+        
         def _jax_wrapped_calc_if_approx(c, a, b, param):
-            return c * a + (1.0 - c) * b
+            sample = c * a + (1.0 - c) * b
+            if debias:
+                hard_sample = jnp.select([c, ~c], [a, b])
+                sample += jax.lax.stop_gradient(hard_sample - sample)
+            return sample
         
         return _jax_wrapped_calc_if_approx, None
     
-    def Switch(self):
+    def control_switch(self):
         if self.verbose:
             raise_warning('Using the replacement rule: '
-                          'switch(pred) { cases } --> sum(cases[i] * (pred == i))')   
+                          'switch(pred) { cases } --> '
+                          'sum(cases[i] * softmax(-abs(pred - i)))')   
             
-        debias = 'Switch' in self.debias
+        debias = 'switch' in self.debias
         
         def _jax_wrapped_calc_switch_approx(pred, cases, param):
-            pred = jnp.broadcast_to(pred[jnp.newaxis, ...], shape=cases.shape)
             literals = FuzzyLogic._literals(cases.shape, axis=0)
+            pred = jnp.broadcast_to(pred[jnp.newaxis, ...], shape=cases.shape)
             proximity = -jnp.abs(pred - literals)
             soft_case = jax.nn.softmax(param * proximity, axis=0)
             sample = jnp.sum(cases * soft_case, axis=0)
@@ -436,8 +503,8 @@ class FuzzyLogic:
     # ===========================================================================
      
     def _gumbel_softmax(self, key, prob):
-        Gumbel01 = random.gumbel(key=key, shape=prob.shape)
-        sample = Gumbel01 + jnp.log1p(prob + self.eps - 1.0)
+        Gumbel01 = random.gumbel(key=key, shape=prob.shape, dtype=self.REAL)
+        sample = Gumbel01 + jnp.log(prob + self.eps)
         return sample
         
     def bernoulli(self):
@@ -448,13 +515,13 @@ class FuzzyLogic:
         jax_gs = self._gumbel_softmax
         jax_argmax, jax_param = self.argmax()
         
-        def _jax_wrapped_calc_switch_approx(key, prob, param):
+        def _jax_wrapped_calc_bernoulli_approx(key, prob, param):
             prob = jnp.stack([1.0 - prob, prob], axis=-1)
             sample = jax_gs(key, prob)
-            sample = jax_argmax(sample, -1, param)
+            sample = jax_argmax(sample, axis=-1, param=param)
             return sample
         
-        return _jax_wrapped_calc_switch_approx, jax_param
+        return _jax_wrapped_calc_bernoulli_approx, jax_param
     
     def discrete(self):
         if self.verbose:
@@ -466,7 +533,7 @@ class FuzzyLogic:
         
         def _jax_wrapped_calc_discrete_approx(key, prob, param):
             sample = jax_gs(key, prob) 
-            sample = jax_argmax(sample, -1, param)
+            sample = jax_argmax(sample, axis=-1, param=param)
             return sample
         
         return _jax_wrapped_calc_discrete_approx, jax_param
@@ -479,11 +546,11 @@ w = 100.0
 
 def _test_logical():
     print('testing logical')
-    _and, _ = logic.And()
-    _not, _ = logic.Not()
+    _and, _ = logic.logical_and()
+    _not, _ = logic.logical_not()
     _gre, _ = logic.greater()
-    _or, _ = logic.Or()
-    _if, _ = logic.If()
+    _or, _ = logic.logical_or()
+    _if, _ = logic.control_if()
     
     # https://towardsdatascience.com/emulating-logical-gates-with-a-neural-network-75c229ec4cc9
     def test_logic(x1, x2):
@@ -516,7 +583,7 @@ def _test_indexing():
 
 def _test_control():
     print('testing control')
-    _switch, _ = logic.Switch()
+    _switch, _ = logic.control_switch()
     
     pred = jnp.asarray(jnp.linspace(0, 2, 10))
     case1 = jnp.asarray([-10.] * 10)