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

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

@@ -1,6 +1,6 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.4
 Name: pyRDDLGym-jax
-Version: 2.3
+Version: 2.5
 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
@@ -39,6 +39,7 @@ Dynamic: description
 Dynamic: description-content-type
 Dynamic: home-page
 Dynamic: license
+Dynamic: license-file
 Dynamic: provides-extra
 Dynamic: requires-dist
 Dynamic: requires-python
@@ -116,7 +117,7 @@ pip install pyRDDLGym-jax[extra,dashboard]
 A basic run script is provided to train JaxPlan on any RDDL problem:
 
 ```shell
-jaxplan plan <domain> <instance> <method> <episodes>
+jaxplan plan <domain> <instance> <method> --episodes <episodes>
 ```
 
 where:
@@ -241,7 +242,7 @@ More documentation about this and other new features will be coming soon.
 A basic run script is provided to run automatic Bayesian hyper-parameter tuning for the most sensitive parameters of JaxPlan:
 
 ```shell
-jaxplan tune <domain> <instance> <method> <trials> <iters> <workers> <dashboard>
+jaxplan tune <domain> <instance> <method> --trials <trials> --iters <iters> --workers <workers> --dashboard <dashboard> --filepath <filepath>
 ```
 
 where:
@@ -251,7 +252,8 @@ where:
 - ``trials`` is the (optional) number of trials/episodes to average in evaluating each hyper-parameter setting
 - ``iters`` is the (optional) maximum number of iterations/evaluations of Bayesian optimization to perform
 - ``workers`` is the (optional) number of parallel evaluations to be done at each iteration, e.g. the total evaluations = ``iters * workers``
-- ``dashboard`` is whether the optimizations are tracked in the dashboard application.
+- ``dashboard`` is whether the optimizations are tracked in the dashboard application
+- ``filepath`` is the optional file path where a config file with the best hyper-parameter setting will be saved.
 
 It is easy to tune a custom range of the planner's hyper-parameters efficiently. 
 First create a config file template with patterns replacing concrete parameter values that you want to tune, e.g.:
@@ -291,23 +293,16 @@ env = pyRDDLGym.make(domain, instance, vectorized=True)
 with open('path/to/config.cfg', 'r') as file: 
     config_template = file.read() 
     
-# map parameters in the config that will be tuned
+# tune weight from 10^-1 ... 10^5 and lr from 10^-5 ... 10^1
 def power_10(x):
-    return 10.0 ** x
-    
-hyperparams = [
-    Hyperparameter('TUNABLE_WEIGHT', -1., 5., power_10),  # tune weight from 10^-1 ... 10^5
-    Hyperparameter('TUNABLE_LEARNING_RATE', -5., 1., power_10),   # tune lr from 10^-5 ... 10^1
-]
+    return 10.0 ** x    
+hyperparams = [Hyperparameter('TUNABLE_WEIGHT', -1., 5., power_10),
+               Hyperparameter('TUNABLE_LEARNING_RATE', -5., 1., power_10)]
     
 # build the tuner and tune
 tuning = JaxParameterTuning(env=env,
-                            config_template=config_template,
-                            hyperparams=hyperparams,
-                            online=False,
-                            eval_trials=trials,
-                            num_workers=workers,
-                            gp_iters=iters)
+                            config_template=config_template, hyperparams=hyperparams,
+                            online=False, eval_trials=trials, num_workers=workers, gp_iters=iters)
 tuning.tune(key=42, log_file='path/to/log.csv')
 ```
   

{pyrddlgym_jax-2.3 → pyrddlgym_jax-2.5}/README.md

@@ -70,7 +70,7 @@ pip install pyRDDLGym-jax[extra,dashboard]
 A basic run script is provided to train JaxPlan on any RDDL problem:
 
 ```shell
-jaxplan plan <domain> <instance> <method> <episodes>
+jaxplan plan <domain> <instance> <method> --episodes <episodes>
 ```
 
 where:
@@ -195,7 +195,7 @@ More documentation about this and other new features will be coming soon.
 A basic run script is provided to run automatic Bayesian hyper-parameter tuning for the most sensitive parameters of JaxPlan:
 
 ```shell
-jaxplan tune <domain> <instance> <method> <trials> <iters> <workers> <dashboard>
+jaxplan tune <domain> <instance> <method> --trials <trials> --iters <iters> --workers <workers> --dashboard <dashboard> --filepath <filepath>
 ```
 
 where:
@@ -205,7 +205,8 @@ where:
 - ``trials`` is the (optional) number of trials/episodes to average in evaluating each hyper-parameter setting
 - ``iters`` is the (optional) maximum number of iterations/evaluations of Bayesian optimization to perform
 - ``workers`` is the (optional) number of parallel evaluations to be done at each iteration, e.g. the total evaluations = ``iters * workers``
-- ``dashboard`` is whether the optimizations are tracked in the dashboard application.
+- ``dashboard`` is whether the optimizations are tracked in the dashboard application
+- ``filepath`` is the optional file path where a config file with the best hyper-parameter setting will be saved.
 
 It is easy to tune a custom range of the planner's hyper-parameters efficiently. 
 First create a config file template with patterns replacing concrete parameter values that you want to tune, e.g.:
@@ -245,23 +246,16 @@ env = pyRDDLGym.make(domain, instance, vectorized=True)
 with open('path/to/config.cfg', 'r') as file: 
     config_template = file.read() 
     
-# map parameters in the config that will be tuned
+# tune weight from 10^-1 ... 10^5 and lr from 10^-5 ... 10^1
 def power_10(x):
-    return 10.0 ** x
-    
-hyperparams = [
-    Hyperparameter('TUNABLE_WEIGHT', -1., 5., power_10),  # tune weight from 10^-1 ... 10^5
-    Hyperparameter('TUNABLE_LEARNING_RATE', -5., 1., power_10),   # tune lr from 10^-5 ... 10^1
-]
+    return 10.0 ** x    
+hyperparams = [Hyperparameter('TUNABLE_WEIGHT', -1., 5., power_10),
+               Hyperparameter('TUNABLE_LEARNING_RATE', -5., 1., power_10)]
     
 # build the tuner and tune
 tuning = JaxParameterTuning(env=env,
-                            config_template=config_template,
-                            hyperparams=hyperparams,
-                            online=False,
-                            eval_trials=trials,
-                            num_workers=workers,
-                            gp_iters=iters)
+                            config_template=config_template, hyperparams=hyperparams,
+                            online=False, eval_trials=trials, num_workers=workers, gp_iters=iters)
 tuning.tune(key=42, log_file='path/to/log.csv')
 ```
   

pyrddlgym_jax-2.5/pyRDDLGym_jax/__init__.py

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

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

@@ -430,7 +430,7 @@ class JaxRDDLCompiler:
                 _jax_wrapped_single_step_policy,
                 in_axes=(0, None, None, None, 0, None)
             )(keys, policy_params, hyperparams, step, subs, model_params)
-            model_params = jax.tree_map(partial(jnp.mean, axis=0), model_params)
+            model_params = jax.tree_util.tree_map(partial(jnp.mean, axis=0), model_params)
             carry = (key, policy_params, hyperparams, subs, model_params)
             return carry, log            
             
@@ -440,7 +440,7 @@ class JaxRDDLCompiler:
             start = (key, policy_params, hyperparams, subs, model_params)
             steps = jnp.arange(n_steps)
             end, log = jax.lax.scan(_jax_wrapped_batched_step_policy, start, steps)
-            log = jax.tree_map(partial(jnp.swapaxes, axis1=0, axis2=1), log)
+            log = jax.tree_util.tree_map(partial(jnp.swapaxes, axis1=0, axis2=1), log)
             model_params = end[-1]
             return log, model_params
         
@@ -471,8 +471,7 @@ class JaxRDDLCompiler:
         return printed
     
     def model_parameter_info(self) -> Dict[str, Dict[str, Any]]:
-        '''Returns a dictionary of additional information about model
-        parameters.'''
+        '''Returns a dictionary of additional information about model parameters.'''
         result = {}
         for (id, value) in self.model_params.items():
             expr_id = int(str(id).split('_')[0])
@@ -708,7 +707,10 @@ class JaxRDDLCompiler:
                     sample = jnp.asarray(value, dtype=self._fix_dtype(value))
                     new_slices = [None] * len(jax_nested_expr)
                     for (i, jax_expr) in enumerate(jax_nested_expr):
-                        new_slices[i], key, err, params = jax_expr(x, params, key)
+                        new_slice, key, err, params = jax_expr(x, params, key)
+                        if not jnp.issubdtype(jnp.result_type(new_slice), jnp.integer):
+                            new_slice = jnp.asarray(new_slice, dtype=self.INT)
+                        new_slices[i] = new_slice
                         error |= err
                     new_slices = tuple(new_slices)
                     sample = sample[new_slices]
@@ -799,7 +801,7 @@ class JaxRDDLCompiler:
         elif n == 2 or (n >= 2 and op in {'*', '+'}):
             jax_exprs = [self._jax(arg, init_params) for arg in args]
             result = jax_exprs[0]
-            for i, jax_rhs in enumerate(jax_exprs[1:]):
+            for (i, jax_rhs) in enumerate(jax_exprs[1:]):
                 jax_op = valid_ops[op](f'{expr.id}_{op}{i}', init_params)
                 result = self._jax_binary(result, jax_rhs, jax_op, at_least_int=True)
             return result
@@ -987,7 +989,8 @@ class JaxRDDLCompiler:
             sample_cases = [None] * len(jax_cases)
             for (i, jax_case) in enumerate(jax_cases):
                 sample_cases[i], key, err_case, params = jax_case(x, params, key)
-                err |= err_case                
+                err |= err_case      
+            sample_cases = jnp.asarray(sample_cases)          
             sample_cases = jnp.asarray(sample_cases, dtype=self._fix_dtype(sample_cases))
             
             # predicate (enum) is an integer - use it to extract from case array

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

@@ -29,6 +29,8 @@
 #
 # ***********************************************************************
 
+
+from abc import ABCMeta, abstractmethod
 import traceback
 from typing import Callable, Dict, Tuple, Union
 
@@ -64,30 +66,35 @@ def enumerate_literals(shape: Tuple[int, ...], axis: int, dtype: type=jnp.int32)
 #
 # ===========================================================================
 
-class Comparison:
+class Comparison(metaclass=ABCMeta):
     '''Base class for approximate comparison operations.'''
     
+    @abstractmethod
     def greater_equal(self, id, init_params):
-        raise NotImplementedError
+        pass
     
+    @abstractmethod
     def greater(self, id, init_params):
-        raise NotImplementedError
+        pass
     
+    @abstractmethod
     def equal(self, id, init_params):
-        raise NotImplementedError
+        pass
     
+    @abstractmethod
     def sgn(self, id, init_params):
-        raise NotImplementedError
+        pass
     
+    @abstractmethod
     def argmax(self, id, init_params):
-        raise NotImplementedError
+        pass
     
 
 class SigmoidComparison(Comparison):
     '''Comparison operations approximated using sigmoid functions.'''
     
     def __init__(self, weight: float=10.0) -> None:
-        self.weight = weight
+        self.weight = float(weight)
         
     # https://arxiv.org/abs/2110.05651
     def greater_equal(self, id, init_params):
@@ -139,21 +146,23 @@ class SigmoidComparison(Comparison):
 #
 # ===========================================================================
 
-class Rounding:
+class Rounding(metaclass=ABCMeta):
     '''Base class for approximate rounding operations.'''
     
+    @abstractmethod
     def floor(self, id, init_params):
-        raise NotImplementedError
+        pass
     
+    @abstractmethod
     def round(self, id, init_params):
-        raise NotImplementedError
+        pass
 
 
 class SoftRounding(Rounding):
     '''Rounding operations approximated using soft operations.'''
     
     def __init__(self, weight: float=10.0) -> None:
-        self.weight = weight
+        self.weight = float(weight)
         
     # https://www.tensorflow.org/probability/api_docs/python/tfp/substrates/jax/bijectors/Softfloor
     def floor(self, id, init_params):
@@ -189,11 +198,12 @@ class SoftRounding(Rounding):
 #
 # ===========================================================================
 
-class Complement:
+class Complement(metaclass=ABCMeta):
     '''Base class for approximate logical complement operations.'''
     
+    @abstractmethod
     def __call__(self, id, init_params):
-        raise NotImplementedError
+        pass
 
 
 class StandardComplement(Complement):
@@ -222,16 +232,18 @@ class StandardComplement(Complement):
 # https://www.sciencedirect.com/science/article/abs/pii/016501149190171L
 # ===========================================================================
 
-class TNorm:
+class TNorm(metaclass=ABCMeta):
     '''Base class for fuzzy differentiable t-norms.'''
     
+    @abstractmethod
     def norm(self, id, init_params):
         '''Elementwise t-norm of x and y.'''
-        raise NotImplementedError
+        pass
     
+    @abstractmethod
     def norms(self, id, init_params):
         '''T-norm computed for tensor x along axis.'''
-        raise NotImplementedError
+        pass
     
 
 class ProductTNorm(TNorm):
@@ -339,26 +351,32 @@ class YagerTNorm(TNorm):
 #
 # ===========================================================================
 
-class RandomSampling:
+class RandomSampling(metaclass=ABCMeta):
     '''Describes how non-reparameterizable random variables are sampled.'''
     
+    @abstractmethod
     def discrete(self, id, init_params, logic):
-        raise NotImplementedError
+        pass
     
+    @abstractmethod
     def poisson(self, id, init_params, logic):
-        raise NotImplementedError
+        pass
     
+    @abstractmethod
     def binomial(self, id, init_params, logic):
-        raise NotImplementedError
+        pass
     
+    @abstractmethod
     def negative_binomial(self, id, init_params, logic):
-        raise NotImplementedError
+        pass
     
+    @abstractmethod
     def geometric(self, id, init_params, logic):
-        raise NotImplementedError
+        pass
     
+    @abstractmethod
     def bernoulli(self, id, init_params, logic):
-        raise NotImplementedError
+        pass
     
     def __str__(self) -> str:
         return 'RandomSampling'
@@ -603,21 +621,23 @@ class Determinization(RandomSampling):
 #
 # ===========================================================================
 
-class ControlFlow:
+class ControlFlow(metaclass=ABCMeta):
     '''A base class for control flow, including if and switch statements.'''
     
+    @abstractmethod
     def if_then_else(self, id, init_params):
-        raise NotImplementedError
+        pass
     
+    @abstractmethod
     def switch(self, id, init_params):
-        raise NotImplementedError
+        pass
 
 
 class SoftControlFlow(ControlFlow):
     '''Soft control flow using a probabilistic interpretation.'''
     
     def __init__(self, weight: float=10.0) -> None:
-        self.weight = weight
+        self.weight = float(weight)
         
     @staticmethod
     def _jax_wrapped_calc_if_then_else_soft(c, a, b, params):
@@ -651,15 +671,15 @@ class SoftControlFlow(ControlFlow):
 # ===========================================================================
 
 
-class Logic:
+class Logic(metaclass=ABCMeta):
     '''A base class for representing logic computations in JAX.'''
     
     def __init__(self, use64bit: bool=False) -> None:
         self.set_use64bit(use64bit)
     
-    def summarize_hyperparameters(self) -> None:
-        print(f'model relaxation:\n'
-              f'    use_64_bit    ={self.use64bit}')
+    def summarize_hyperparameters(self) -> str:
+        return (f'model relaxation:\n'
+                f'    use_64_bit    ={self.use64bit}')
     
     def set_use64bit(self, use64bit: bool) -> None:
         '''Toggles whether or not the JAX system will use 64 bit precision.'''
@@ -765,119 +785,150 @@ class Logic:
     # ===========================================================================
     # logical operators
     # ===========================================================================
-     
+    
+    @abstractmethod
     def logical_and(self, id, init_params):
-        raise NotImplementedError
+        pass
     
+    @abstractmethod
     def logical_not(self, id, init_params):
-        raise NotImplementedError
+        pass
     
+    @abstractmethod
     def logical_or(self, id, init_params):
-        raise NotImplementedError
+        pass
     
+    @abstractmethod
     def xor(self, id, init_params):
-        raise NotImplementedError
+        pass
     
+    @abstractmethod
     def implies(self, id, init_params):
-        raise NotImplementedError
+        pass
     
+    @abstractmethod
     def equiv(self, id, init_params):
-        raise NotImplementedError
+        pass
     
+    @abstractmethod
     def forall(self, id, init_params):
-        raise NotImplementedError
+        pass
     
+    @abstractmethod
     def exists(self, id, init_params):    
-        raise NotImplementedError 
+        pass
     
     # ===========================================================================
     # comparison operators
     # ===========================================================================
     
+    @abstractmethod
     def greater_equal(self, id, init_params):
-        raise NotImplementedError 
+        pass
     
+    @abstractmethod
     def greater(self, id, init_params):
-        raise NotImplementedError 
+        pass
     
+    @abstractmethod
     def less_equal(self, id, init_params):
-        raise NotImplementedError 
+        pass
     
+    @abstractmethod
     def less(self, id, init_params):
-        raise NotImplementedError 
+        pass
     
+    @abstractmethod
     def equal(self, id, init_params):
-        raise NotImplementedError 
+        pass
     
+    @abstractmethod
     def not_equal(self, id, init_params):
-        raise NotImplementedError 
+        pass
     
     # ===========================================================================
     # special functions
     # ===========================================================================
      
+    @abstractmethod
     def sgn(self, id, init_params):
-        raise NotImplementedError 
+        pass
     
+    @abstractmethod
     def floor(self, id, init_params):
-        raise NotImplementedError 
+        pass
     
+    @abstractmethod
     def round(self, id, init_params):
-        raise NotImplementedError 
+        pass
     
+    @abstractmethod
     def ceil(self, id, init_params):
-        raise NotImplementedError 
+        pass 
     
+    @abstractmethod
     def div(self, id, init_params):
-        raise NotImplementedError 
+        pass 
     
+    @abstractmethod
     def mod(self, id, init_params):
-        raise NotImplementedError 
+        pass
     
+    @abstractmethod
     def sqrt(self, id, init_params):
-        raise NotImplementedError 
+        pass
     
     # ===========================================================================
     # indexing
     # ===========================================================================
-     
+    
+    @abstractmethod
     def argmax(self, id, init_params):   
-        raise NotImplementedError  
+        pass
     
+    @abstractmethod
     def argmin(self, id, init_params):   
-        raise NotImplementedError     
+        pass
     
     # ===========================================================================
     # control flow
     # ===========================================================================
      
+    @abstractmethod
     def control_if(self, id, init_params):
-        raise NotImplementedError
+        pass
         
+    @abstractmethod
     def control_switch(self, id, init_params):
-        raise NotImplementedError
+        pass
     
     # ===========================================================================
     # random variables
     # ===========================================================================
      
+    @abstractmethod
     def discrete(self, id, init_params):
-        raise NotImplementedError
+        pass
     
+    @abstractmethod
     def bernoulli(self, id, init_params):
-        raise NotImplementedError
+        pass
     
+    @abstractmethod
     def poisson(self, id, init_params):
-        raise NotImplementedError
+        pass
     
+    @abstractmethod
     def geometric(self, id, init_params):
-        raise NotImplementedError
+        pass
     
+    @abstractmethod
     def binomial(self, id, init_params):
-        raise NotImplementedError
+        pass
 
+    @abstractmethod
     def negative_binomial(self, id, init_params):
-        raise NotImplementedError
+        pass
 
 
 class ExactLogic(Logic):
@@ -1109,8 +1160,8 @@ class FuzzyLogic(Logic):
                 f'    underflow_tol={self.eps}\n'
                 f'    use_64_bit   ={self.use64bit}\n')
 
-    def summarize_hyperparameters(self) -> None:
-        print(self.__str__())
+    def summarize_hyperparameters(self) -> str:
+        return self.__str__()
         
     # ===========================================================================
     # logical operators