pyRDDLGym-jax 0.3__py3-none-any.whl → 0.5__py3-none-any.whl

pyRDDLGym_jax/__init__.py

@@ -1 +1 @@
-__version__ = '0.3'
+__version__ = '0.4'

pyRDDLGym_jax/core/compiler.py

@@ -1,22 +1,11 @@
 from functools import partial
+import traceback
+from typing import Any, Callable, Dict, List, Optional
+
 import jax
 import jax.numpy as jnp
 import jax.random as random
 import jax.scipy as scipy 
-import traceback
-from typing import Any, Callable, Dict, List, Optional
-
-from pyRDDLGym.core.debug.exception import raise_warning
-
-# 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 complex distributions '
-                  '(Binomial, Negative-Binomial, Multinomial) will fail.', 'red')
-    traceback.print_exc()
-    tfp = None
 
 from pyRDDLGym.core.compiler.initializer import RDDLValueInitializer
 from pyRDDLGym.core.compiler.levels import RDDLLevelAnalysis
@@ -25,12 +14,23 @@ from pyRDDLGym.core.compiler.tracer import RDDLObjectsTracer
 from pyRDDLGym.core.constraints import RDDLConstraints
 from pyRDDLGym.core.debug.exception import (
     print_stack_trace, 
+    raise_warning,
     RDDLInvalidNumberOfArgumentsError, 
     RDDLNotImplementedError
 )
 from pyRDDLGym.core.debug.logger import Logger
 from pyRDDLGym.core.simulator import RDDLSimulatorPrecompiled
 
+# 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 complex distributions '
+                  '(Binomial, Negative-Binomial, Multinomial) will fail.', 'red')
+    traceback.print_exc()
+    tfp = None
+
 
 # ===========================================================================
 # EXACT RDDL TO JAX COMPILATION RULES
@@ -87,7 +87,7 @@ def _function_aggregation_exact_named(op, name):
 def _function_if_exact_named():
         
     def _jax_wrapped_if_exact(c, a, b, param):
-        return jnp.where(c, a, b)
+        return jnp.where(c > 0.5, a, b)
         
     return _jax_wrapped_if_exact
     
@@ -114,16 +114,27 @@ def _function_bernoulli_exact_named():
 def _function_discrete_exact_named():
         
     def _jax_wrapped_discrete_exact(key, prob, param):
-        logits = jnp.log(prob)
-        sample = random.categorical(key=key, logits=logits, axis=-1)
-        out_of_bounds = jnp.logical_not(jnp.logical_and(
-            jnp.all(prob >= 0),
-            jnp.allclose(jnp.sum(prob, axis=-1), 1.0)))
-        return sample, out_of_bounds
+        return random.categorical(key=key, logits=jnp.log(prob), axis=-1)
         
     return _jax_wrapped_discrete_exact
 
 
+def _function_poisson_exact_named():
+    
+    def _jax_wrapped_poisson_exact(key, rate, param):
+        return random.poisson(key=key, lam=rate, dtype=jnp.int64)
+    
+    return _jax_wrapped_poisson_exact
+
+
+def _function_geometric_exact_named():
+    
+    def _jax_wrapped_geometric_exact(key, prob, param):
+        return random.geometric(key=key, p=prob, dtype=jnp.int64)
+    
+    return _jax_wrapped_geometric_exact
+
+
 class JaxRDDLCompiler:
     '''Compiles a RDDL AST representation into an equivalent JAX representation.
     All operations are identical to their numpy equivalents.
@@ -210,12 +221,12 @@ class JaxRDDLCompiler:
     }
     
     EXACT_RDDL_TO_JAX_IF = _function_if_exact_named()
-    
     EXACT_RDDL_TO_JAX_SWITCH = _function_switch_exact_named()
     
     EXACT_RDDL_TO_JAX_BERNOULLI = _function_bernoulli_exact_named()
-    
     EXACT_RDDL_TO_JAX_DISCRETE = _function_discrete_exact_named()
+    EXACT_RDDL_TO_JAX_POISSON = _function_poisson_exact_named()
+    EXACT_RDDL_TO_JAX_GEOMETRIC = _function_geometric_exact_named()
 
     def __init__(self, rddl: RDDLLiftedModel,
                  allow_synchronous_state: bool=True,
@@ -289,6 +300,8 @@ class JaxRDDLCompiler:
         self.SWITCH_HELPER = JaxRDDLCompiler.EXACT_RDDL_TO_JAX_SWITCH
         self.BERNOULLI_HELPER = JaxRDDLCompiler.EXACT_RDDL_TO_JAX_BERNOULLI
         self.DISCRETE_HELPER = JaxRDDLCompiler.EXACT_RDDL_TO_JAX_DISCRETE
+        self.POISSON_HELPER = JaxRDDLCompiler.EXACT_RDDL_TO_JAX_POISSON
+        self.GEOMETRIC_HELPER = JaxRDDLCompiler.EXACT_RDDL_TO_JAX_GEOMETRIC
     
     # ===========================================================================
     # main compilation subroutines
@@ -996,13 +1009,14 @@ class JaxRDDLCompiler:
             jax_op, jax_param = self._unwrap(negative_op, expr.id, info)
             return self._jax_unary(jax_expr, jax_op, jax_param, at_least_int=True)
                     
-        elif n == 2:
-            lhs, rhs = args
-            jax_lhs = self._jax(lhs, info)
-            jax_rhs = self._jax(rhs, info)
+        elif n == 2 or (n >= 2 and op in {'*', '+'}):
+            jax_exprs = [self._jax(arg, info) for arg in args]
             jax_op, jax_param = self._unwrap(valid_ops[op], expr.id, info)
-            return self._jax_binary(
-                jax_lhs, jax_rhs, jax_op, jax_param, at_least_int=True)
+            result = jax_exprs[0]
+            for jax_rhs in jax_exprs[1:]:
+                result = self._jax_binary(
+                    result, jax_rhs, jax_op, jax_param, at_least_int=True)
+            return result
         
         JaxRDDLCompiler._check_num_args(expr, 2)
     
@@ -1046,13 +1060,14 @@ class JaxRDDLCompiler:
             jax_op, jax_param = self._unwrap(logical_not_op, expr.id, info)
             return self._jax_unary(jax_expr, jax_op, jax_param, check_dtype=bool)
         
-        elif n == 2:
-            lhs, rhs = args
-            jax_lhs = self._jax(lhs, info)
-            jax_rhs = self._jax(rhs, info)
+        elif n == 2 or (n >= 2 and op in {'^', '&', '|'}):
+            jax_exprs = [self._jax(arg, info) for arg in args]
             jax_op, jax_param = self._unwrap(valid_ops[op], expr.id, info)
-            return self._jax_binary(
-                jax_lhs, jax_rhs, jax_op, jax_param, check_dtype=bool)
+            result = jax_exprs[0]
+            for jax_rhs in jax_exprs[1:]:
+                result = self._jax_binary(
+                    result, jax_rhs, jax_op, jax_param, check_dtype=bool)
+            return result
         
         JaxRDDLCompiler._check_num_args(expr, 2)
     
@@ -1165,16 +1180,17 @@ class JaxRDDLCompiler:
         return _jax_wrapped_if_then_else
     
     def _jax_switch(self, expr, info):
+        pred, *_ = expr.args
              
-        # if expression is non-fluent, always use the exact operation
-        if self.compile_non_fluent_exact and not self.traced.cached_is_fluent(expr):
+        # if predicate is non-fluent, always use the exact operation
+        # case conditions are currently only literals so they are non-fluent
+        if self.compile_non_fluent_exact and not self.traced.cached_is_fluent(pred):
             switch_op = JaxRDDLCompiler.EXACT_RDDL_TO_JAX_SWITCH
         else:
             switch_op = self.SWITCH_HELPER
         jax_switch, jax_param = self._unwrap(switch_op, expr.id, info)
         
         # recursively compile predicate
-        pred, *_ = expr.args  
         jax_pred = self._jax(pred, info)
         
         # recursively compile cases
@@ -1426,15 +1442,24 @@ class JaxRDDLCompiler:
     def _jax_poisson(self, expr, info):
         ERR = JaxRDDLCompiler.ERROR_CODES['INVALID_PARAM_POISSON']
         JaxRDDLCompiler._check_num_args(expr, 1)
-        
         arg_rate, = expr.args
+        
+        # if rate is non-fluent, always use the exact operation
+        if self.compile_non_fluent_exact and not self.traced.cached_is_fluent(arg_rate):
+            poisson_op = JaxRDDLCompiler.EXACT_RDDL_TO_JAX_POISSON
+        else:
+            poisson_op = self.POISSON_HELPER
+        jax_poisson, jax_param = self._unwrap(poisson_op, expr.id, info)
+        
+        # recursively compile arguments
         jax_rate = self._jax(arg_rate, info)
         
         # uses the implicit JAX subroutine
         def _jax_wrapped_distribution_poisson(x, params, key):
             rate, key, err = jax_rate(x, params, key)
             key, subkey = random.split(key)
-            sample = random.poisson(key=subkey, lam=rate, dtype=self.INT)
+            param = params.get(jax_param, None)
+            sample = jax_poisson(subkey, rate, param).astype(self.INT)
             out_of_bounds = jnp.logical_not(jnp.all(rate >= 0))
             err |= (out_of_bounds * ERR)
             return sample, key, err
@@ -1535,33 +1560,25 @@ class JaxRDDLCompiler:
         ERR = JaxRDDLCompiler.ERROR_CODES['INVALID_PARAM_GEOMETRIC']
         JaxRDDLCompiler._check_num_args(expr, 1)        
         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_prob):
+            geom_op = JaxRDDLCompiler.EXACT_RDDL_TO_JAX_GEOMETRIC
+        else:
+            geom_op = self.GEOMETRIC_HELPER
+        jax_geom, jax_param = self._unwrap(geom_op, expr.id, info)
+        
+        # recursively compile arguments        
         jax_prob = self._jax(arg_prob, info)
         
-        if self.compile_non_fluent_exact and not self.traced.cached_is_fluent(arg_prob):            
-            
-            # prob is non-fluent: do not reparameterize
-            def _jax_wrapped_distribution_geometric(x, params, key):
-                prob, key, err = jax_prob(x, params, key)
-                key, subkey = random.split(key)
-                sample = random.geometric(key=subkey, p=prob, dtype=self.INT)
-                out_of_bounds = jnp.logical_not(jnp.all((prob >= 0) & (prob <= 1)))
-                err |= (out_of_bounds * ERR)
-                return sample, key, err
-        
-        else:    
-            floor_op, jax_param = self._unwrap(
-                self.KNOWN_UNARY['floor'], expr.id, info)
-            
-            # reparameterization trick Geom(p) = floor(ln(U(0, 1)) / ln(p)) + 1
-            def _jax_wrapped_distribution_geometric(x, params, key):
-                prob, key, err = jax_prob(x, params, key)
-                key, subkey = random.split(key)
-                U = random.uniform(key=subkey, shape=jnp.shape(prob), dtype=self.REAL)
-                param = params.get(jax_param, None)
-                sample = floor_op(jnp.log(U) / jnp.log(1.0 - prob), param) + 1
-                out_of_bounds = jnp.logical_not(jnp.all((prob >= 0) & (prob <= 1)))
-                err |= (out_of_bounds * ERR)
-                return sample, key, err
+        def _jax_wrapped_distribution_geometric(x, params, key):
+            prob, key, err = jax_prob(x, params, key)
+            key, subkey = random.split(key)
+            param = params.get(jax_param, None)
+            sample = jax_geom(subkey, prob, param).astype(self.INT)
+            out_of_bounds = jnp.logical_not(jnp.all((prob >= 0) & (prob <= 1)))
+            err |= (out_of_bounds * ERR)
+            return sample, key, err
         
         return _jax_wrapped_distribution_geometric
     
@@ -1770,7 +1787,10 @@ class JaxRDDLCompiler:
             # dispatch to sampling subroutine
             key, subkey = random.split(key)
             param = params.get(jax_param, None)
-            sample, out_of_bounds = jax_discrete(subkey, prob, param)
+            sample = jax_discrete(subkey, prob, param)
+            out_of_bounds = jnp.logical_not(jnp.logical_and(
+                jnp.all(prob >= 0),
+                jnp.allclose(jnp.sum(prob, axis=-1), 1.0)))
             error |= (out_of_bounds * ERR)
             return sample, key, error
         
@@ -1803,7 +1823,10 @@ class JaxRDDLCompiler:
             # dispatch to sampling subroutine
             key, subkey = random.split(key)
             param = params.get(jax_param, None)
-            sample, out_of_bounds = jax_discrete(subkey, prob, param)
+            sample = jax_discrete(subkey, prob, param)
+            out_of_bounds = jnp.logical_not(jnp.logical_and(
+                jnp.all(prob >= 0),
+                jnp.allclose(jnp.sum(prob, axis=-1), 1.0)))
             error |= (out_of_bounds * ERR)
             return sample, key, error