pyRDDLGym-jax 2.5__py3-none-any.whl → 2.6__py3-none-any.whl

pyRDDLGym_jax/__init__.py

@@ -1 +1 @@
-__version__ = '2.5'
+__version__ = '2.6'

pyRDDLGym_jax/core/compiler.py

@@ -237,7 +237,8 @@ class JaxRDDLCompiler:
     
     def compile_transition(self, check_constraints: bool=False,
                            constraint_func: bool=False, 
-                           init_params_constr: Dict[str, Any]={}) -> Callable:
+                           init_params_constr: Dict[str, Any]={},
+                           cache_path_info: bool=False) -> Callable:
         '''Compiles the current RDDL model into a JAX transition function that 
         samples the next state.
         
@@ -274,6 +275,7 @@ class JaxRDDLCompiler:
         returned log and does not raise an exception
         :param constraint_func: produces the h(s, a) function described above
         in addition to the usual outputs
+        :param cache_path_info: whether to save full path traces as part of the log
         '''
         NORMAL = JaxRDDLCompiler.ERROR_CODES['NORMAL']        
         rddl = self.rddl
@@ -322,8 +324,11 @@ class JaxRDDLCompiler:
             errors |= err
             
             # calculate fluent values
-            fluents = {name: values for (name, values) in subs.items() 
-                       if name not in rddl.non_fluents}
+            if cache_path_info:
+                fluents = {name: values for (name, values) in subs.items() 
+                           if name not in rddl.non_fluents}
+            else:
+                fluents = {}
             
             # set the next state to the current state
             for (state, next_state) in rddl.next_state.items():
@@ -367,7 +372,9 @@ class JaxRDDLCompiler:
                          n_batch: int,
                          check_constraints: bool=False,
                          constraint_func: bool=False, 
-                         init_params_constr: Dict[str, Any]={}) -> Callable:
+                         init_params_constr: Dict[str, Any]={},
+                         model_params_reduction: Callable=lambda x: x[0],
+                         cache_path_info: bool=False) -> Callable:
         '''Compiles the current RDDL model into a JAX transition function that 
         samples trajectories with a fixed horizon from a policy.
         
@@ -399,10 +406,13 @@ class JaxRDDLCompiler:
         returned log and does not raise an exception
         :param constraint_func: produces the h(s, a) constraint function
         in addition to the usual outputs
+        :param model_params_reduction: how to aggregate updated model_params across runs
+        in the batch (defaults to selecting the first element's parameters in the batch)
+        :param cache_path_info: whether to save full path traces as part of the log
         '''
         rddl = self.rddl
         jax_step_fn = self.compile_transition(
-            check_constraints, constraint_func, init_params_constr)
+            check_constraints, constraint_func, init_params_constr, cache_path_info)
         
         # for POMDP only observ-fluents are assumed visible to the policy
         if rddl.observ_fluents:
@@ -421,7 +431,6 @@ class JaxRDDLCompiler:
             return jax_step_fn(subkey, actions, subs, model_params)
                         
         # do a batched step update from the policy
-        # TODO: come up with a better way to reduce the model_param batch dim
         def _jax_wrapped_batched_step_policy(carry, step):
             key, policy_params, hyperparams, subs, model_params = carry  
             key, *subkeys = random.split(key, num=1 + n_batch)
@@ -430,7 +439,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_util.tree_map(partial(jnp.mean, axis=0), model_params)
+            model_params = jax.tree_util.tree_map(model_params_reduction, model_params)
             carry = (key, policy_params, hyperparams, subs, model_params)
             return carry, log            
             

pyRDDLGym_jax/core/logic.py

@@ -1056,15 +1056,13 @@ class ExactLogic(Logic):
     def control_if(self, id, init_params):
         return self._jax_wrapped_calc_if_then_else_exact
     
-    @staticmethod
-    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._jax_wrapped_calc_switch_exact
+        def _jax_wrapped_calc_switch_exact(pred, cases, params):
+            pred = jnp.asarray(pred[jnp.newaxis, ...], dtype=self.INT)
+            sample = jnp.take_along_axis(cases, pred, axis=0)
+            assert sample.shape[0] == 1
+            return sample[0, ...], params
+        return _jax_wrapped_calc_switch_exact
     
     # ===========================================================================
     # random variables

pyRDDLGym_jax/core/model.py

@@ -0,0 +1,595 @@
+from collections import deque
+from copy import deepcopy
+from enum import Enum
+from functools import partial
+import sys
+import time
+from tqdm import tqdm
+from typing import Any, Callable, Dict, Generator, Iterable, Optional, Tuple
+
+import jax
+import jax.nn.initializers as initializers
+import jax.numpy as jnp
+import jax.random as random
+import numpy as np
+import optax
+
+from pyRDDLGym.core.compiler.model import RDDLLiftedModel
+
+from pyRDDLGym_jax.core.logic import Logic, ExactLogic
+from pyRDDLGym_jax.core.planner import JaxRDDLCompilerWithGrad
+
+Kwargs = Dict[str, Any]
+State = Dict[str, np.ndarray]
+Action = Dict[str, np.ndarray]
+DataStream = Iterable[Tuple[State, Action, State]]
+Params = Dict[str, np.ndarray]
+Callback = Dict[str, Any]
+LossFunction = Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]
+
+
+# ***********************************************************************
+# ALL VERSIONS OF LOSS FUNCTIONS
+#
+# - loss functions based on specific likelihood assumptions (MSE, cross-entropy)
+# 
+# ***********************************************************************
+
+
+def mean_squared_error() -> LossFunction:
+    def _jax_wrapped_mse_loss(target, pred):
+        loss_values = jnp.square(target - pred)
+        return loss_values
+    return jax.jit(_jax_wrapped_mse_loss)
+
+
+def binary_cross_entropy(eps: float=1e-6) -> LossFunction:
+    def _jax_wrapped_binary_cross_entropy_loss(target, pred):
+        pred = jnp.clip(pred, eps, 1.0 - eps)
+        log_pred = jnp.log(pred)
+        log_not_pred = jnp.log(1.0 - pred)
+        loss_values = -target * log_pred - (1.0 - target) * log_not_pred
+        return loss_values
+    return jax.jit(_jax_wrapped_binary_cross_entropy_loss)
+
+
+def optax_loss(loss_fn: LossFunction, **kwargs) -> LossFunction:
+    def _jax_wrapped_optax_loss(target, pred):
+        loss_values = loss_fn(pred, target, **kwargs)
+        return loss_values
+    return jax.jit(_jax_wrapped_optax_loss)
+
+
+# ***********************************************************************
+# ALL VERSIONS OF JAX MODEL LEARNER
+#
+# - gradient based model learning
+# 
+# ***********************************************************************
+
+
+class JaxLearnerStatus(Enum):
+    '''Represents the status of a parameter update from the JAX model learner, 
+    including whether the update resulted in nan gradient, 
+    whether progress was made, budget was reached, or other information that
+    can be used to monitor and act based on the learner's progress.'''
+    
+    NORMAL = 0
+    NO_PROGRESS = 1
+    INVALID_GRADIENT = 2
+    TIME_BUDGET_REACHED = 3
+    ITER_BUDGET_REACHED = 4
+    
+    def is_terminal(self) -> bool:
+        return self.value >= 2
+    
+
+class JaxModelLearner:
+    '''A class for data-driven estimation of unknown parameters in a given RDDL MDP using 
+    gradient descent.'''
+
+    def __init__(self, rddl: RDDLLiftedModel, 
+                 param_ranges: Dict[str, Tuple[Optional[np.ndarray], Optional[np.ndarray]]],
+                 batch_size_train: int=32,
+                 samples_per_datapoint: int=1,
+                 optimizer: Callable[..., optax.GradientTransformation]=optax.rmsprop,
+                 optimizer_kwargs: Optional[Kwargs]=None,
+                 initializer: initializers.Initializer = initializers.normal(),
+                 wrap_non_bool: bool=True,
+                 use64bit: bool=False,
+                 bool_fluent_loss: LossFunction=binary_cross_entropy(),
+                 real_fluent_loss: LossFunction=mean_squared_error(),
+                 int_fluent_loss: LossFunction=mean_squared_error(),
+                 logic: Logic=ExactLogic(),
+                 model_params_reduction: Callable=lambda x: x[0]) -> None:
+        '''Creates a new gradient-based algorithm for inferring unknown non-fluents
+        in a RDDL domain from a data set or stream coming from the real environment.
+        
+        :param rddl: the RDDL domain to learn
+        :param param_ranges: the ranges of all learnable non-fluents
+        :param batch_size_train: how many transitions to compute per optimization 
+        step
+        :param samples_per_datapoint: how many random samples to produce from the step
+        function per data point during training
+        :param optimizer: a factory for an optax SGD algorithm
+        :param optimizer_kwargs: a dictionary of parameters to pass to the SGD
+        factory (e.g. which parameters are controllable externally)
+        :param initializer: how to initialize non-fluents
+        :param wrap_non_bool: whether to wrap non-boolean trainable parameters to satisfy
+        required ranges as specified in param_ranges (use a projected gradient otherwise)
+        :param use64bit: whether to perform arithmetic in 64 bit
+        :param bool_fluent_loss: loss function to optimize for bool-valued fluents
+        :param real_fluent_loss: loss function to optimize for real-valued fluents
+        :param int_fluent_loss: loss function to optimize for int-valued fluents
+        :param logic: a subclass of Logic for mapping exact mathematical
+        operations to their differentiable counterparts 
+        :param model_params_reduction: how to aggregate updated model_params across runs
+        in the batch (defaults to selecting the first element's parameters in the batch)
+        '''
+        self.rddl = rddl
+        self.param_ranges = param_ranges.copy()
+        self.batch_size_train = batch_size_train
+        self.samples_per_datapoint = samples_per_datapoint
+        if optimizer_kwargs is None:
+            optimizer_kwargs = {'learning_rate': 0.001}
+        self.optimizer_kwargs = optimizer_kwargs
+        self.initializer = initializer
+        self.wrap_non_bool = wrap_non_bool
+        self.use64bit = use64bit
+        self.bool_fluent_loss = bool_fluent_loss
+        self.real_fluent_loss = real_fluent_loss
+        self.int_fluent_loss = int_fluent_loss
+        self.logic = logic
+        self.model_params_reduction = model_params_reduction
+
+        # validate param_ranges
+        for (name, values) in param_ranges.items():
+            if name not in rddl.non_fluents:
+                raise ValueError(
+                    f'param_ranges key <{name}> is not a valid non-fluent '
+                    f'in the current rddl.')
+            if not isinstance(values, (tuple, list)):
+                raise ValueError(
+                    f'param_ranges values with key <{name}> are neither a tuple nor a list.')
+            if len(values) != 2:
+                raise ValueError(
+                    f'param_ranges values with key <{name}> must be of length 2, '
+                    f'got length {len(values)}.')
+            lower, upper = values
+            if lower is not None and upper is not None and not np.all(lower <= upper):
+                raise ValueError(
+                    f'param_ranges values with key <{name}> do not satisfy lower <= upper.')
+
+        # build the optimizer
+        optimizer = optimizer(**optimizer_kwargs)
+        pipeline = [optimizer]
+        self.optimizer = optax.chain(*pipeline)
+
+        # build the computation graph
+        self.step_fn = self._jax_compile_rddl()
+        self.map_fn = self._jax_map()
+        self.loss_fn = self._jax_loss(map_fn=self.map_fn, step_fn=self.step_fn)
+        self.update_fn, self.project_fn = self._jax_update(loss_fn=self.loss_fn)
+        self.init_fn, self.init_opt_fn = self._jax_init(project_fn=self.project_fn)
+    
+    # ===========================================================================
+    # COMPILATION SUBROUTINES
+    # ===========================================================================
+
+    def _jax_compile_rddl(self):
+
+        # compile the RDDL model
+        self.compiled = JaxRDDLCompilerWithGrad(
+            rddl=self.rddl,
+            logic=self.logic,
+            use64bit=self.use64bit,
+            compile_non_fluent_exact=False,
+            print_warnings=True
+        )
+        self.compiled.compile(log_jax_expr=True, heading='RELAXED MODEL')
+
+        # compile the transition step function
+        step_fn = self.compiled.compile_transition()
+
+        def _jax_wrapped_step(key, param_fluents, subs, actions, hyperparams):
+            for (name, param) in param_fluents.items():
+                subs[name] = param
+            subs, _, hyperparams = step_fn(key, actions, subs, hyperparams)
+            return subs, hyperparams
+
+        # batched step function
+        def _jax_wrapped_batched_step(key, param_fluents, subs, actions, hyperparams):
+            keys = jnp.asarray(random.split(key, num=self.batch_size_train))
+            subs, hyperparams = jax.vmap(
+                _jax_wrapped_step, in_axes=(0, None, 0, 0, None)
+            )(keys, param_fluents, subs, actions, hyperparams)
+            hyperparams = jax.tree_util.tree_map(self.model_params_reduction, hyperparams)
+            return subs, hyperparams
+
+        # batched step function with parallel samples per data point
+        def _jax_wrapped_batched_parallel_step(key, param_fluents, subs, actions, hyperparams):
+            keys = jnp.asarray(random.split(key, num=self.samples_per_datapoint))
+            subs, hyperparams = jax.vmap(
+                _jax_wrapped_batched_step, in_axes=(0, None, None, None, None)
+            )(keys, param_fluents, subs, actions, hyperparams)
+            hyperparams = jax.tree_util.tree_map(self.model_params_reduction, hyperparams)
+            return subs, hyperparams
+
+        batched_step_fn = jax.jit(_jax_wrapped_batched_parallel_step)
+        return batched_step_fn        
+
+    def _jax_map(self):
+
+        # compute case indices for bounding
+        case_indices = {}
+        if self.wrap_non_bool:
+            for (name, (lower, upper)) in self.param_ranges.items():
+                if lower is None: lower = -np.inf
+                if upper is None: upper = +np.inf
+                self.param_ranges[name] = (lower, upper)
+                case_indices[name] = (
+                    0 * (np.isfinite(lower) & np.isfinite(upper)) +
+                    1 * (np.isfinite(lower) & ~np.isfinite(upper)) +
+                    2 * (~np.isfinite(lower) & np.isfinite(upper)) +
+                    3 * (~np.isfinite(lower) & ~np.isfinite(upper))
+                )
+        
+        # map trainable parameters to their non-fluent values
+        def _jax_wrapped_params_to_fluents(params):
+            param_fluents = {}
+            for (name, param) in params.items():
+                if self.rddl.variable_ranges[name] == 'bool':
+                    param_fluents[name] = jax.nn.sigmoid(param)
+                else:
+                    if self.wrap_non_bool:
+                        lower, upper = self.param_ranges[name]
+                        cases = [
+                            lambda x: lower + (upper - lower) * jax.nn.sigmoid(x),
+                            lambda x: lower + (jax.nn.elu(x) + 1.0),
+                            lambda x: upper - (jax.nn.elu(-x) + 1.0),
+                            lambda x: x
+                        ]
+                        indices = case_indices[name]
+                        param_fluents[name] = jax.lax.switch(indices, cases, param)
+                    else:
+                        param_fluents[name] = param
+            return param_fluents
+        
+        map_fn = jax.jit(_jax_wrapped_params_to_fluents)
+        return map_fn
+
+    def _jax_loss(self, map_fn, step_fn):
+
+        # use binary cross entropy for bool fluents
+        # mean squared error for continuous and integer fluents
+        def _jax_wrapped_batched_model_loss(key, param_fluents, subs, actions, next_fluents, 
+                                            hyperparams):
+            next_subs, hyperparams = step_fn(key, param_fluents, subs, actions, hyperparams)
+            total_loss = 0.0
+            for (name, next_value) in next_fluents.items():
+                preds = jnp.asarray(next_subs[name], dtype=self.compiled.REAL)
+                targets = jnp.asarray(next_value, dtype=self.compiled.REAL)[jnp.newaxis, ...]
+                if self.rddl.variable_ranges[name] == 'bool':
+                    loss_values = self.bool_fluent_loss(targets, preds)
+                elif self.rddl.variable_ranges[name] == 'real':
+                    loss_values = self.real_fluent_loss(targets, preds)
+                else:
+                    loss_values = self.int_fluent_loss(targets, preds)
+                total_loss += jnp.mean(loss_values) / len(next_fluents)
+            return total_loss, hyperparams
+        
+        # loss with the parameters mapped to their fluents
+        def _jax_wrapped_batched_loss(key, params, subs, actions, next_fluents, hyperparams):
+            param_fluents = map_fn(params)
+            loss, hyperparams = _jax_wrapped_batched_model_loss(
+                key, param_fluents, subs, actions, next_fluents, hyperparams)
+            return loss, hyperparams
+
+        loss_fn = jax.jit(_jax_wrapped_batched_loss)
+        return loss_fn
+    
+    def _jax_init(self, project_fn):
+        optimizer = self.optimizer
+        
+        # initialize both the non-fluents and optimizer
+        def _jax_wrapped_init_params_optimizer(key):
+            params = {}
+            for name in self.param_ranges:
+                shape = jnp.shape(self.compiled.init_values[name])
+                key, subkey = random.split(key)
+                params[name] = self.initializer(subkey, shape, dtype=self.compiled.REAL)
+            params = project_fn(params)
+            opt_state = optimizer.init(params)
+            return params, opt_state
+        
+        # initialize just the optimizer given the non-fluents
+        def _jax_wrapped_init_optimizer(params):
+            params = project_fn(params)
+            opt_state = optimizer.init(params)
+            return params, opt_state
+        
+        init_fn = jax.jit(_jax_wrapped_init_params_optimizer)
+        init_opt_fn = jax.jit(_jax_wrapped_init_optimizer)
+        return init_fn, init_opt_fn
+                    
+    def _jax_update(self, loss_fn):
+        optimizer = self.optimizer
+
+        # projected gradient trick to satisfy box constraints on params
+        def _jax_wrapped_project_params(params):
+            if self.wrap_non_bool:
+                return params
+            else:
+                new_params = {}
+                for (name, value) in params.items():
+                    if self.rddl.variable_ranges[name] == 'bool':
+                        new_params[name] = value
+                    else:
+                        lower, upper = self.param_ranges[name]
+                        new_params[name] = jnp.clip(value, lower, upper)
+                return new_params
+
+        # gradient descent update
+        def _jax_wrapped_params_update(key, params, subs, actions, next_fluents, 
+                                       hyperparams, opt_state):
+            (loss_val, hyperparams), grad = jax.value_and_grad(
+                loss_fn, argnums=1, has_aux=True
+            )(key, params, subs, actions, next_fluents, hyperparams)
+            updates, opt_state = optimizer.update(grad, opt_state)
+            params = optax.apply_updates(params, updates)
+            params = _jax_wrapped_project_params(params)
+            zero_grads = jax.tree_util.tree_map(partial(jnp.allclose, b=0.0), grad)
+            return params, opt_state, loss_val, zero_grads, hyperparams
+    
+        update_fn = jax.jit(_jax_wrapped_params_update)
+        project_fn = jax.jit(_jax_wrapped_project_params)
+        return update_fn, project_fn
+    
+    def _batched_init_subs(self): 
+        init_train = {}
+        for (name, value) in self.compiled.init_values.items():
+            value = np.reshape(value, np.shape(value))[np.newaxis, ...]
+            value = np.repeat(value, repeats=self.batch_size_train, axis=0)
+            value = np.asarray(value, dtype=self.compiled.REAL)
+            init_train[name] = value
+        for (state, next_state) in self.rddl.next_state.items():
+            init_train[next_state] = init_train[state]
+        return init_train
+    
+    # ===========================================================================
+    # ESTIMATE API
+    # ===========================================================================
+
+    def optimize(self, *args, **kwargs) -> Optional[Callback]:
+        '''Estimate the unknown parameters from the given data set. 
+        Return the callback from training.
+
+        :param data: a data stream represented as a (possibly infinite) sequence of 
+        transition batches of the form (states, actions, next-states), where each element
+        is a numpy array of leading dimension equal to batch_size_train
+        :param key: JAX PRNG key (derived from clock if not provided)
+        :param epochs: the maximum number of steps of gradient descent
+        :param train_seconds: total time allocated for gradient descent  
+        :param guess: initial non-fluent parameters: if None will use the initializer
+        specified in this instance   
+        :param print_progress: whether to print the progress bar during training
+        '''
+        it = self.optimize_generator(*args, **kwargs)
+        
+        # https://stackoverflow.com/questions/50937966/fastest-most-pythonic-way-to-consume-an-iterator
+        callback = None
+        if sys.implementation.name == 'cpython':
+            last_callback = deque(it, maxlen=1)
+            if last_callback:
+                callback = last_callback.pop()
+        else:
+            for callback in it:
+                pass
+        return callback
+
+    def optimize_generator(self, data: DataStream,
+                           key: Optional[random.PRNGKey]=None,
+                           epochs: int=999999,
+                           train_seconds: float=120.,
+                           guess: Optional[Params]=None,
+                           print_progress: bool=True) -> Generator[Callback, None, None]:
+        '''Return a generator for estimating the unknown parameters from the given data set. 
+        Generator can be iterated over to lazily estimate the parameters, yielding
+        a dictionary of intermediate computations.
+
+        :param data: a data stream represented as a (possibly infinite) sequence of 
+        transition batches of the form (states, actions, next-states), where each element
+        is a numpy array of leading dimension equal to batch_size_train
+        :param key: JAX PRNG key (derived from clock if not provided)
+        :param epochs: the maximum number of steps of gradient descent
+        :param train_seconds: total time allocated for gradient descent
+        :param guess: initial non-fluent parameters: if None will use the initializer
+        specified in this instance
+        :param print_progress: whether to print the progress bar during training
+        '''
+        start_time = time.time()
+        elapsed_outside_loop = 0
+
+        # if PRNG key is not provided
+        if key is None:
+            key = random.PRNGKey(round(time.time() * 1000))
+        
+        # prepare initial subs
+        subs = self._batched_init_subs()
+
+        # initialize parameter fluents to optimize
+        if guess is None:
+            key, subkey = random.split(key)
+            params, opt_state = self.init_fn(subkey)
+        else:
+            params, opt_state = self.init_opt_fn(guess)
+
+        # initialize model hyper-parameters
+        hyperparams = self.compiled.model_params
+
+        # progress bar
+        if print_progress:
+            progress_bar = tqdm(
+                None, total=100, bar_format='{l_bar}{bar}| {elapsed} {postfix}')
+        else:
+            progress_bar = None
+
+        # main training loop
+        for (it, (states, actions, next_states)) in enumerate(data):
+            status = JaxLearnerStatus.NORMAL
+
+            # gradient update
+            subs.update(states)
+            key, subkey = random.split(key)
+            params, opt_state, loss, zero_grads, hyperparams = self.update_fn(
+                subkey, params, subs, actions, next_states, hyperparams, opt_state)
+            
+            # extract non-fluent values from the trainable parameters
+            param_fluents = self.map_fn(params)
+            param_fluents = {name: param_fluents[name] for name in self.param_ranges}
+
+            # check for learnability
+            params_zero_grads = {
+                name for (name, zero_grad) in zero_grads.items() if zero_grad}
+            if params_zero_grads:
+                status = JaxLearnerStatus.NO_PROGRESS
+
+            # reached computation budget
+            elapsed = time.time() - start_time - elapsed_outside_loop
+            if elapsed >= train_seconds:
+                status = JaxLearnerStatus.TIME_BUDGET_REACHED
+            if it >= epochs - 1:
+                status = JaxLearnerStatus.ITER_BUDGET_REACHED
+            
+            # build a callback
+            progress_percent = 100 * min(
+                1, max(0, elapsed / train_seconds, it / (epochs - 1)))
+            callback = {
+                'status': status,
+                'iteration': it,
+                'train_loss': loss,
+                'params': params,
+                'param_fluents': param_fluents,
+                'key': key,
+                'progress': progress_percent
+            }
+            
+            # update progress
+            if print_progress:
+                progress_bar.set_description(
+                    f'{it:7} it / {loss:12.8f} train / {status.value} status', refresh=False)
+                progress_bar.set_postfix_str(
+                    f'{(it + 1) / (elapsed + 1e-6):.2f}it/s', refresh=False)
+                progress_bar.update(progress_percent - progress_bar.n)
+            
+            # yield the callback
+            start_time_outside = time.time()
+            yield callback
+            elapsed_outside_loop += (time.time() - start_time_outside)
+
+            # abortion check
+            if status.is_terminal():
+                break
+    
+    def evaluate_loss(self, data: DataStream,
+                      key: Optional[random.PRNGKey], 
+                      param_fluents: Params) -> float:
+        '''Evaluates the model loss of the given learned non-fluent values and the data.
+
+        :param data: a data stream represented as a (possibly infinite) sequence of 
+        transition batches of the form (states, actions, next-states), where each element
+        is a numpy array of leading dimension equal to batch_size_train
+        :param key: JAX PRNG key (derived from clock if not provided)
+        :param param_fluents: the learned non-fluent values
+        '''
+        if key is None:
+            key = random.PRNGKey(round(time.time() * 1000))
+        subs = self._batched_init_subs()
+        hyperparams = self.compiled.model_params
+        mean_loss = 0.0
+        for (it, (states, actions, next_states)) in enumerate(data):
+            subs.update(states)
+            key, subkey = random.split(key)
+            loss_value, _ = self.loss_fn(
+                subkey, param_fluents, subs, actions, next_states, hyperparams)
+            mean_loss += (loss_value - mean_loss) / (it + 1)
+        return mean_loss
+
+    def learned_model(self, param_fluents: Params) -> RDDLLiftedModel:
+        '''Substitutes the given learned non-fluent values into the RDDL model and returns
+        the new model.
+
+        :param param_fluents: the learned non-fluent values
+        '''
+        model = deepcopy(self.rddl)
+        for (name, values) in param_fluents.items():
+            prange = model.variable_ranges[name]
+            if prange == 'real':
+                pass
+            elif prange == 'bool':
+                values = values > 0.5
+            else:
+                values = np.asarray(values, dtype=self.compiled.INT)
+            values = np.ravel(values, order='C').tolist()
+            if not self.rddl.variable_params[name]:
+                assert(len(values) == 1)
+                values = values[0]
+            model.non_fluents[name] = values
+        return model
+
+
+if __name__ == '__main__':
+    import os
+    import pyRDDLGym
+    from pyRDDLGym_jax.core.planner import load_config, JaxBackpropPlanner, JaxOfflineController
+    bs = 32
+
+    # make some data
+    def data_iterator():
+        env = pyRDDLGym.make('CartPole_Continuous_gym', '0', vectorized=True)
+        model = JaxModelLearner(rddl=env.model, param_ranges={}, batch_size_train=bs)
+        key = random.PRNGKey(round(time.time() * 1000))
+        subs = model._batched_init_subs()
+        param_fluents = {}
+        while True:
+            states = {
+                'pos': np.random.uniform(-2.4, 2.4, (bs,)),
+                'vel': np.random.uniform(-2.4, 2.4, (bs,)),
+                'ang-pos': np.random.uniform(-0.21, 0.21, (bs,)),
+                'ang-vel': np.random.uniform(-0.21, 0.21, (bs,))
+            }
+            subs.update(states)
+            actions = {
+                'force': np.random.uniform(-10., 10., (bs,))
+            }
+            key, subkey = random.split(key)
+            subs, _ = model.step_fn(subkey, param_fluents, subs, actions, {})
+            subs = {k: np.asarray(v)[0, ...] for k, v in subs.items()}
+            next_states = {k: subs[k] for k in model.rddl.state_fluents}
+            yield (states, actions, next_states)
+    
+    # train it
+    env = pyRDDLGym.make('TestJax', '0', vectorized=True)
+    model_learner = JaxModelLearner(rddl=env.model, 
+                                    param_ranges={
+                                        'w1': (None, None), 'b1': (None, None),
+                                        'w2': (None, None), 'b2': (None, None),
+                                        'w1o': (None, None), 'b1o': (None, None),
+                                        'w2o': (None, None), 'b2o': (None, None)
+                                    }, 
+                                    batch_size_train=bs, 
+                                    optimizer_kwargs = {'learning_rate': 0.0003})
+    for cb in model_learner.optimize_generator(data_iterator(), epochs=10000):
+        pass
+
+    # planning in the trained model
+    model = model_learner.learned_model(cb['param_fluents'])
+    abs_path = os.path.dirname(os.path.abspath(__file__))        
+    config_path = os.path.join(os.path.dirname(abs_path), 'examples', 'configs', 'default_drp.cfg') 
+    planner_args, _, train_args = load_config(config_path)
+    planner = JaxBackpropPlanner(rddl=model, **planner_args)
+    controller = JaxOfflineController(planner, **train_args)
+
+    # evaluation of the plan
+    test_env = pyRDDLGym.make('CartPole_Continuous_gym', '0', vectorized=True)
+    controller.evaluate(test_env, episodes=1, verbose=True, render=True)

pyRDDLGym_jax/core/planner.py

@@ -207,6 +207,13 @@ def _load_config(config, args):
                 pgpe_kwargs['optimizer'] = pgpe_optimizer
         planner_args['pgpe'] = getattr(sys.modules[__name__], pgpe_method)(**pgpe_kwargs)
 
+    # preprocessor settings
+    preproc_method = planner_args.get('preprocessor', None)
+    preproc_kwargs = planner_args.pop('preprocessor_kwargs', {})
+    if preproc_method is not None:
+        planner_args['preprocessor'] = getattr(
+            sys.modules[__name__], preproc_method)(**preproc_kwargs)
+
     # optimize call RNG key
     planner_key = train_args.get('key', None)
     if planner_key is not None:
@@ -343,6 +350,100 @@ class JaxRDDLCompilerWithGrad(JaxRDDLCompiler):
         return arg
 
 
+# ***********************************************************************
+# ALL VERSIONS OF STATE PREPROCESSING FOR DRP
+# 
+# - static normalization
+# 
+# ***********************************************************************
+
+
+class Preprocessor(metaclass=ABCMeta):
+    '''Base class for all state preprocessors.'''
+
+    HYPERPARAMS_KEY = 'preprocessor__'
+
+    def __init__(self) -> None:
+        self._initializer = None
+        self._update = None
+        self._transform = None
+
+    @property
+    def initialize(self):
+        return self._initializer
+
+    @property
+    def update(self):
+        return self._update
+    
+    @property
+    def transform(self):
+        return self._transform
+    
+    @abstractmethod
+    def compile(self, compiled: JaxRDDLCompilerWithGrad) -> None:
+        pass
+
+
+class StaticNormalizer(Preprocessor):
+    '''Normalize values by box constraints on fluents computed from the RDDL domain.'''
+
+    def __init__(self, fluent_bounds: Dict[str, Tuple[np.ndarray, np.ndarray]]={}) -> None:
+        '''Create a new instance of the static normalizer.
+
+        :param fluent_bounds: optional bounds on fluents to overwrite default values.
+        '''
+        self.fluent_bounds = fluent_bounds
+
+    def compile(self, compiled: JaxRDDLCompilerWithGrad) -> None:
+        
+        # adjust for partial observability
+        rddl = compiled.rddl
+        if rddl.observ_fluents:
+            observed_vars = rddl.observ_fluents
+        else:
+            observed_vars = rddl.state_fluents
+
+        # ignore boolean fluents and infinite bounds
+        bounded_vars = {}
+        for var in observed_vars:
+            if rddl.variable_ranges[var] != 'bool':
+                lower, upper = compiled.constraints.bounds[var]
+                if np.all(np.isfinite(lower) & np.isfinite(upper) & (lower < upper)):
+                    bounded_vars[var] = (lower, upper)
+                user_bounds = self.fluent_bounds.get(var, None)
+                if user_bounds is not None:
+                    bounded_vars[var] = tuple(user_bounds)
+        
+        # initialize to ranges computed by the constraint parser
+        def _jax_wrapped_normalizer_init():
+            return bounded_vars        
+        self._initializer = jax.jit(_jax_wrapped_normalizer_init)
+
+        # static bounds
+        def _jax_wrapped_normalizer_update(subs, stats):
+            stats = {var: (jnp.asarray(lower, dtype=compiled.REAL), 
+                           jnp.asarray(upper, dtype=compiled.REAL)) 
+                     for (var, (lower, upper)) in bounded_vars.items()}     
+            return stats   
+        self._update = jax.jit(_jax_wrapped_normalizer_update)
+
+        # apply min max scaling
+        def _jax_wrapped_normalizer_transform(subs, stats):
+            new_subs = {}
+            for (var, values) in subs.items():
+                if var in stats:
+                    lower, upper = stats[var]
+                    new_dims = jnp.ndim(values) - jnp.ndim(lower)
+                    lower = lower[(jnp.newaxis,) * new_dims + (...,)]
+                    upper = upper[(jnp.newaxis,) * new_dims + (...,)]
+                    new_subs[var] = (values - lower) / (upper - lower)
+                else:
+                    new_subs[var] = values
+            return new_subs
+        self._transform = jax.jit(_jax_wrapped_normalizer_transform)
+
+
 # ***********************************************************************
 # ALL VERSIONS OF JAX PLANS
 # 
@@ -368,7 +469,8 @@ class JaxPlan(metaclass=ABCMeta):
     @abstractmethod
     def compile(self, compiled: JaxRDDLCompilerWithGrad,
                 _bounds: Bounds,
-                horizon: int) -> None:
+                horizon: int,
+                preprocessor: Optional[Preprocessor]=None) -> None:
         pass
     
     @abstractmethod
@@ -519,7 +621,8 @@ class JaxStraightLinePlan(JaxPlan):
     
     def compile(self, compiled: JaxRDDLCompilerWithGrad,
                 _bounds: Bounds,
-                horizon: int) -> None:
+                horizon: int,
+                preprocessor: Optional[Preprocessor]=None) -> None:
         rddl = compiled.rddl
         
         # calculate the correct action box bounds
@@ -607,7 +710,7 @@ class JaxStraightLinePlan(JaxPlan):
             return new_params, True
         
         # convert softmax action back to action dict
-        action_sizes = {var: np.prod(shape[1:], dtype=int) 
+        action_sizes = {var: np.prod(shape[1:], dtype=np.int64) 
                         for (var, shape) in shapes.items()
                         if ranges[var] == 'bool'}
         
@@ -691,7 +794,7 @@ class JaxStraightLinePlan(JaxPlan):
                 scores = []
                 for (var, param) in params.items():
                     if ranges[var] == 'bool':
-                        param_flat = jnp.ravel(param)
+                        param_flat = jnp.ravel(param, order='C')
                         if noop[var]:
                             if wrap_sigmoid:
                                 param_flat = -param_flat
@@ -908,7 +1011,8 @@ class JaxDeepReactivePolicy(JaxPlan):
         
     def compile(self, compiled: JaxRDDLCompilerWithGrad,
                 _bounds: Bounds,
-                horizon: int) -> None:
+                horizon: int,
+                preprocessor: Optional[Preprocessor]=None) -> None:
         rddl = compiled.rddl
         
         # calculate the correct action box bounds
@@ -939,7 +1043,7 @@ class JaxDeepReactivePolicy(JaxPlan):
         wrap_non_bool = self._wrap_non_bool
         init = self._initializer
         layers = list(enumerate(zip(self._topology, self._activations)))
-        layer_sizes = {var: np.prod(shape, dtype=int) 
+        layer_sizes = {var: np.prod(shape, dtype=np.int64) 
                        for (var, shape) in shapes.items()}
         layer_names = {var: f'output_{var}'.replace('-', '_') for var in shapes}
         
@@ -973,7 +1077,12 @@ class JaxDeepReactivePolicy(JaxPlan):
                 normalize = False
         
         # convert subs dictionary into a state vector to feed to the MLP
-        def _jax_wrapped_policy_input(subs):
+        def _jax_wrapped_policy_input(subs, hyperparams):
+
+            # optional state preprocessing
+            if preprocessor is not None:
+                stats = hyperparams[preprocessor.HYPERPARAMS_KEY]
+                subs = preprocessor.transform(subs, stats)
             
             # concatenate all state variables into a single vector
             # optionally apply layer norm to each input tensor
@@ -981,7 +1090,7 @@ class JaxDeepReactivePolicy(JaxPlan):
             non_bool_dims = 0
             for (var, value) in subs.items():
                 if var in observed_vars:
-                    state = jnp.ravel(value)
+                    state = jnp.ravel(value, order='C')
                     if ranges[var] == 'bool':
                         states_bool.append(state)
                     else:
@@ -1010,8 +1119,8 @@ class JaxDeepReactivePolicy(JaxPlan):
             return state
             
         # predict actions from the policy network for current state
-        def _jax_wrapped_policy_network_predict(subs):
-            state = _jax_wrapped_policy_input(subs)
+        def _jax_wrapped_policy_network_predict(subs, hyperparams):
+            state = _jax_wrapped_policy_input(subs, hyperparams)
             
             # feed state vector through hidden layers
             hidden = state
@@ -1076,7 +1185,7 @@ class JaxDeepReactivePolicy(JaxPlan):
         
         # train action prediction
         def _jax_wrapped_drp_predict_train(key, params, hyperparams, step, subs):
-            actions = predict_fn.apply(params, subs)
+            actions = predict_fn.apply(params, subs, hyperparams)
             if not wrap_non_bool:
                 for (var, action) in actions.items():
                     if var != bool_key and ranges[var] != 'bool':
@@ -1126,7 +1235,7 @@ class JaxDeepReactivePolicy(JaxPlan):
             subs = {var: value[0, ...] 
                     for (var, value) in subs.items()
                     if var in observed_vars}
-            params = predict_fn.init(key, subs)
+            params = predict_fn.init(key, subs, hyperparams)
             return params
         
         self.initializer = _jax_wrapped_drp_init
@@ -1634,12 +1743,21 @@ def mean_semivariance_utility(returns: jnp.ndarray, beta: float) -> float:
     return mu - 0.5 * beta * msv
 
 
+@jax.jit
+def sharpe_utility(returns: jnp.ndarray, risk_free: float) -> float:
+    return (jnp.mean(returns) - risk_free) / (jnp.std(returns) + 1e-10)
+
+
+@jax.jit
+def var_utility(returns: jnp.ndarray, alpha: float) -> float:
+    return jnp.percentile(returns, q=100 * alpha)
+
+
 @jax.jit
 def cvar_utility(returns: jnp.ndarray, alpha: float) -> float:
     var = jnp.percentile(returns, q=100 * alpha)
     mask = returns <= var
-    weights = mask / jnp.maximum(1, jnp.sum(mask))
-    return jnp.sum(returns * weights)
+    return jnp.sum(returns * mask) / jnp.maximum(1, jnp.sum(mask))
 
 
 # set of all currently valid built-in utility functions
@@ -1649,8 +1767,10 @@ UTILITY_LOOKUP = {
     'mean_std': mean_deviation_utility,
     'mean_semivar': mean_semivariance_utility,
     'mean_semidev': mean_semideviation_utility,
+    'sharpe': sharpe_utility,
     'entropic': entropic_utility,
     'exponential': entropic_utility,
+    'var': var_utility,
     'cvar': cvar_utility
 }
 
@@ -1689,7 +1809,8 @@ class JaxBackpropPlanner:
                  logger: Optional[Logger]=None,
                  dashboard_viz: Optional[Any]=None,
                  print_warnings: bool=True,
-                 parallel_updates: Optional[int]=None) -> None:
+                 parallel_updates: Optional[int]=None,
+                 preprocessor: Optional[Preprocessor]=None) -> None:
         '''Creates a new gradient-based algorithm for optimizing action sequences
         (plan) in the given RDDL. Some operations will be converted to their
         differentiable counterparts; the specific operations can be customized
@@ -1731,6 +1852,7 @@ class JaxBackpropPlanner:
         to pass to the dashboard to visualize the policy
         :param print_warnings: whether to print warnings
         :param parallel_updates: how many optimizers to run independently in parallel
+        :param preprocessor: optional preprocessor for state inputs to plan
         '''
         self.rddl = rddl
         self.plan = plan
@@ -1756,6 +1878,7 @@ class JaxBackpropPlanner:
         self.pgpe = pgpe
         self.use_pgpe = pgpe is not None
         self.print_warnings = print_warnings
+        self.preprocessor = preprocessor
         
         # set optimizer
         try:
@@ -1881,7 +2004,8 @@ r"""
                   f'    noise_kwargs      ={self.noise_kwargs}\n'
                   f'    batch_size_train  ={self.batch_size_train}\n'
                   f'    batch_size_test   ={self.batch_size_test}\n'
-                  f'    parallel_updates  ={self.parallel_updates}\n')
+                  f'    parallel_updates  ={self.parallel_updates}\n'
+                  f'    preprocessor      ={self.preprocessor}\n')
         result += str(self.plan)
         if self.use_pgpe:
             result += str(self.pgpe)
@@ -1917,10 +2041,15 @@ r"""
         
     def _jax_compile_optimizer(self):
         
+        # preprocessor
+        if self.preprocessor is not None:
+            self.preprocessor.compile(self.compiled)   
+
         # policy
         self.plan.compile(self.compiled,
                           _bounds=self._action_bounds,
-                          horizon=self.horizon)
+                          horizon=self.horizon,
+                          preprocessor=self.preprocessor)
         self.train_policy = jax.jit(self.plan.train_policy)
         self.test_policy = jax.jit(self.plan.test_policy)
         
@@ -1928,14 +2057,16 @@ r"""
         train_rollouts = self.compiled.compile_rollouts(
             policy=self.plan.train_policy,
             n_steps=self.horizon,
-            n_batch=self.batch_size_train
+            n_batch=self.batch_size_train,
+            cache_path_info=self.preprocessor is not None
         )
         self.train_rollouts = train_rollouts
         
         test_rollouts = self.test_compiled.compile_rollouts(
             policy=self.plan.test_policy,
             n_steps=self.horizon,
-            n_batch=self.batch_size_test
+            n_batch=self.batch_size_test,
+            cache_path_info=False
         )
         self.test_rollouts = jax.jit(test_rollouts)
         
@@ -2397,7 +2528,13 @@ r"""
                             f'which could be suboptimal.', 'yellow')
                         print(message)
                     policy_hyperparams[action] = 1.0
-            
+        
+        # initialize preprocessor
+        preproc_key = None
+        if self.preprocessor is not None:
+            preproc_key = self.preprocessor.HYPERPARAMS_KEY
+            policy_hyperparams[preproc_key] = self.preprocessor.initialize()
+
         # print summary of parameters:
         if print_summary:
             print(self.summarize_system())
@@ -2524,6 +2661,11 @@ r"""
                  subkey, policy_params, policy_hyperparams, train_subs, model_params, 
                  opt_state, opt_aux)
             
+            # update the preprocessor
+            if self.preprocessor is not None:                
+                policy_hyperparams[preproc_key] = self.preprocessor.update(
+                    train_log['fluents'], policy_hyperparams[preproc_key])
+
             # evaluate
             test_loss, (test_log, model_params_test) = self.test_loss(
                 subkey, policy_params, policy_hyperparams, test_subs, model_params_test)
@@ -2676,6 +2818,7 @@ r"""
                 'model_params': model_params,
                 'progress': progress_percent,
                 'train_log': train_log,
+                'policy_hyperparams': policy_hyperparams,
                 **test_log
             }
 
@@ -2753,7 +2896,8 @@ r"""
     
     def _perform_diagnosis(self, last_iter_improve,
                            train_return, test_return, best_return, grad_norm):
-        max_grad_norm = max(jax.tree_util.tree_leaves(grad_norm))
+        grad_norms = jax.tree_util.tree_leaves(grad_norm)
+        max_grad_norm = max(grad_norms) if grad_norms else np.nan
         grad_is_zero = np.allclose(max_grad_norm, 0)
         
         # divergence if the solution is not finite
@@ -2895,6 +3039,7 @@ class JaxOfflineController(BaseAgent):
         self.train_on_reset = train_on_reset
         self.train_kwargs = train_kwargs        
         self.params_given = params is not None
+        self.hyperparams_given = eval_hyperparams is not None
         
         # load the policy from file
         if not self.train_on_reset and params is not None and isinstance(params, str):
@@ -2908,6 +3053,8 @@ class JaxOfflineController(BaseAgent):
             callback = self.planner.optimize(key=self.key, **self.train_kwargs)
             self.callback = callback
             params = callback['best_params'] 
+            if not self.hyperparams_given:
+                self.eval_hyperparams = callback['policy_hyperparams']
 
             # save the policy
             if save_path is not None:
@@ -2931,6 +3078,8 @@ class JaxOfflineController(BaseAgent):
             callback = self.planner.optimize(key=self.key, **self.train_kwargs)
             self.callback = callback
             self.params = callback['best_params']
+            if not self.hyperparams_given:
+                self.eval_hyperparams = callback['policy_hyperparams']
 
 
 class JaxOnlineController(BaseAgent):
@@ -2963,6 +3112,7 @@ class JaxOnlineController(BaseAgent):
             key = random.PRNGKey(round(time.time() * 1000))
         self.key = key
         self.eval_hyperparams = eval_hyperparams
+        self.hyperparams_given = eval_hyperparams is not None
         self.warm_start = warm_start
         self.train_kwargs = train_kwargs
         self.max_attempts = max_attempts
@@ -2987,6 +3137,8 @@ class JaxOnlineController(BaseAgent):
                 key=self.key, guess=self.guess, subs=state, **self.train_kwargs)    
         self.callback = callback
         params = callback['best_params']
+        if not self.hyperparams_given:
+            self.eval_hyperparams = callback['policy_hyperparams']
 
         # get the action from the parameters for the current state
         self.key, subkey = random.split(self.key)

{pyrddlgym_jax-2.5.dist-info → pyrddlgym_jax-2.6.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pyRDDLGym-jax
-Version: 2.5
+Version: 2.6
 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
@@ -20,7 +20,7 @@ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Requires-Python: >=3.9
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: pyRDDLGym>=2.0
+Requires-Dist: pyRDDLGym>=2.3
 Requires-Dist: tqdm>=4.66
 Requires-Dist: jax>=0.4.12
 Requires-Dist: optax>=0.1.9
@@ -55,7 +55,7 @@ Dynamic: summary
 
 [Installation](#installation) | [Run cmd](#running-from-the-command-line) | [Run python](#running-from-another-python-application) | [Configuration](#configuring-the-planner) | [Dashboard](#jaxplan-dashboard) | [Tuning](#tuning-the-planner) | [Simulation](#simulation) | [Citing](#citing-jaxplan)
 
-**pyRDDLGym-jax (known in the literature as JaxPlan) is an efficient gradient-based/differentiable planning algorithm in JAX.** 
+**pyRDDLGym-jax (or JaxPlan) is an efficient gradient-based planning algorithm based on JAX.** 
 
 Purpose:
 
@@ -84,7 +84,7 @@ and was moved to the individual logic components which have their own unique wei
 
 > [!NOTE]  
 > While JaxPlan can support some discrete state/action problems through model relaxations, on some discrete problems it can perform poorly (though there is an ongoing effort to remedy this!).
-> If you find it is not making sufficient progress, check out the [PROST planner](https://github.com/pyrddlgym-project/pyRDDLGym-prost) (for discrete spaces) or the [deep reinforcement learning wrappers](https://github.com/pyrddlgym-project/pyRDDLGym-rl).
+> If you find it is not making progress, check out the [PROST planner](https://github.com/pyrddlgym-project/pyRDDLGym-prost) (for discrete spaces) or the [deep reinforcement learning wrappers](https://github.com/pyrddlgym-project/pyRDDLGym-rl).
   
 ## Installation
 
@@ -220,13 +220,7 @@ controller = JaxOfflineController(planner, **train_args)
 ## JaxPlan Dashboard
 
 Since version 1.0, JaxPlan has an optional dashboard that allows keeping track of the planner performance across multiple runs, 
-and visualization of the policy or model, and other useful debugging features.
-
-<p align="middle">
-<img src="https://github.com/pyrddlgym-project/pyRDDLGym-jax/blob/main/Images/dashboard.png" width="480" height="248" margin=0/>
-</p>
-
-To run the dashboard, add the following entry to your config file:
+and visualization of the policy or model, and other useful debugging features. To run the dashboard, add the following to your config file:
 
 ```ini
 ...
@@ -235,8 +229,6 @@ dashboard=True
 ...
 ```
 
-More documentation about this and other new features will be coming soon.
-
 ## Tuning the Planner
 
 A basic run script is provided to run automatic Bayesian hyper-parameter tuning for the most sensitive parameters of JaxPlan:

{pyrddlgym_jax-2.5.dist-info → pyrddlgym_jax-2.6.dist-info}/RECORD

@@ -1,9 +1,10 @@
-pyRDDLGym_jax/__init__.py,sha256=VoxLo_sy8RlJIIyu7szqL-cdMGBJdQPg-aSeyOVVIkY,19
+pyRDDLGym_jax/__init__.py,sha256=VUmQViJtwUg1JGcgXlmNm0fE3Njyruyt_76c16R-LTo,19
 pyRDDLGym_jax/entry_point.py,sha256=K0zy1oe66jfBHkHHCM6aGHbbiVqnQvDhDb8se4uaKHE,3319
 pyRDDLGym_jax/core/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-pyRDDLGym_jax/core/compiler.py,sha256=uFCtoipsIa3MM9nGgT3X8iCViPl2XSPNXh0jMdzN0ko,82895
-pyRDDLGym_jax/core/logic.py,sha256=lfc2ak_ap_ajMEFlB5EHCRNgJym31dNyA-5d-7N4CZA,56271
-pyRDDLGym_jax/core/planner.py,sha256=M6GKzN7Ml57B4ZrFZhhkpsQCvReKaCQNzer7zeHCM9E,140275
+pyRDDLGym_jax/core/compiler.py,sha256=Bpgfw4nqRFqiTju7ioR0B0Dhp3wMvk-9LmTRpMmLIOc,83457
+pyRDDLGym_jax/core/logic.py,sha256=9rRpKJCx4Us_2c6BiSWRN9k2sM_iYsAK1B7zcgwu3ZA,56290
+pyRDDLGym_jax/core/model.py,sha256=4WfmtUVN1EKCD-7eWeQByWk8_zKyDcMABAMdlxN1LOU,27215
+pyRDDLGym_jax/core/planner.py,sha256=a684ss5TAkJ-P2SEbZA90FSpDwFxHwRoaLtbRIBspAA,146450
 pyRDDLGym_jax/core/simulator.py,sha256=ayCATTUL3clLaZPQ5OUg2bI_c26KKCTq6TbrxbMsVdc,10470
 pyRDDLGym_jax/core/tuning.py,sha256=BWcQZk02TMLexTz1Sw4lX2EQKvmPbp7biC51M-IiNUw,25153
 pyRDDLGym_jax/core/visualization.py,sha256=4BghMp8N7qtF0tdyDSqtxAxNfP9HPrQWTiXzAMJmx7o,70365
@@ -41,9 +42,9 @@ pyRDDLGym_jax/examples/configs/default_slp.cfg,sha256=mJo0woDevhQCSQfJg30ULVy9qG
 pyRDDLGym_jax/examples/configs/tuning_drp.cfg,sha256=zocZn_cVarH5i0hOlt2Zu0NwmXYBmTTghLaXLtQOGto,526
 pyRDDLGym_jax/examples/configs/tuning_replan.cfg,sha256=9oIhtw9cuikmlbDgCgbrTc5G7hUio-HeAv_3CEGVclY,523
 pyRDDLGym_jax/examples/configs/tuning_slp.cfg,sha256=QqnyR__5-HhKeCDfGDel8VIlqsjxRHk4SSH089zJP8s,486
-pyrddlgym_jax-2.5.dist-info/licenses/LICENSE,sha256=Y0Gi6H6mLOKN-oIKGZulQkoTJyPZeAaeuZu7FXH-meg,1095
-pyrddlgym_jax-2.5.dist-info/METADATA,sha256=XAaEJfbsYW-txxZhFZ6o_HmvqxkIMTqBF9LbV-KdTzI,17058
-pyrddlgym_jax-2.5.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-pyrddlgym_jax-2.5.dist-info/entry_points.txt,sha256=Q--z9QzqDBz1xjswPZ87PU-pib-WPXx44hUWAFoBGBA,59
-pyrddlgym_jax-2.5.dist-info/top_level.txt,sha256=n_oWkP_BoZK0VofvPKKmBZ3NPk86WFNvLhi1BktCbVQ,14
-pyrddlgym_jax-2.5.dist-info/RECORD,,
+pyrddlgym_jax-2.6.dist-info/licenses/LICENSE,sha256=Y0Gi6H6mLOKN-oIKGZulQkoTJyPZeAaeuZu7FXH-meg,1095
+pyrddlgym_jax-2.6.dist-info/METADATA,sha256=1gY3EPRHKMVeZYYgq4DCqWvw3Q1Ak5XVYRaIO2UlQXc,16770
+pyrddlgym_jax-2.6.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+pyrddlgym_jax-2.6.dist-info/entry_points.txt,sha256=Q--z9QzqDBz1xjswPZ87PU-pib-WPXx44hUWAFoBGBA,59
+pyrddlgym_jax-2.6.dist-info/top_level.txt,sha256=n_oWkP_BoZK0VofvPKKmBZ3NPk86WFNvLhi1BktCbVQ,14
+pyrddlgym_jax-2.6.dist-info/RECORD,,