pyRDDLGym-jax 1.0__tar.gz → 1.1__tar.gz

{pyrddlgym_jax-1.0 → pyrddlgym_jax-1.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pyRDDLGym-jax
-Version: 1.0
+Version: 1.1
 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
@@ -31,7 +31,17 @@ Requires-Dist: dash-bootstrap-components>=1.6.0; extra == "dashboard"
 
 # pyRDDLGym-jax
 
-**pyRDDLGym-jax (known in the literature as JaxPlan) is an efficient gradient-based/differentiable planning algorithm in JAX.** It provides:
+![Python Version](https://img.shields.io/badge/python-3.9%2B-blue)
+[![PyPI Version](https://img.shields.io/pypi/v/pyRDDLGym-jax.svg)](https://pypi.org/project/pyRDDLGym-jax/)
+[![Documentation Status](https://readthedocs.org/projects/pyrddlgym/badge/?version=latest)](https://pyrddlgym.readthedocs.io/en/latest/jax.html)
+![License: MIT](https://img.shields.io/badge/License-MIT-blue.svg)
+[![Cumulative PyPI Downloads](https://img.shields.io/pypi/dm/pyrddlgym-jax)](https://pypistats.org/packages/pyrddlgym-jax)
+
+[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.** 
+
+Purpose:
 
 1. automatic translation of any RDDL description file into a differentiable simulator in JAX
 2. flexible policy class representations, automatic model relaxations for working in discrete and hybrid domains, and Bayesian hyper-parameter tuning.
@@ -56,17 +66,6 @@ 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).
-
-## Contents
-
-- [Installation](#installation)
-- [Running from the Command Line](#running-from-the-command-line)
-- [Running from Another Python Application](#running-from-another-python-application)
-- [Configuring the Planner](#configuring-the-planner)
-- [JaxPlan Dashboard](#jaxplan-dashboard)
-- [Tuning the Planner](#tuning-the-planner)
-- [Simulation](#simulation)
-- [Citing JaxPlan](#citing-jaxplan)
   
 ## Installation
 

{pyrddlgym_jax-1.0 → pyrddlgym_jax-1.1}/README.md

@@ -1,6 +1,16 @@
 # pyRDDLGym-jax
 
-**pyRDDLGym-jax (known in the literature as JaxPlan) is an efficient gradient-based/differentiable planning algorithm in JAX.** It provides:
+![Python Version](https://img.shields.io/badge/python-3.9%2B-blue)
+[![PyPI Version](https://img.shields.io/pypi/v/pyRDDLGym-jax.svg)](https://pypi.org/project/pyRDDLGym-jax/)
+[![Documentation Status](https://readthedocs.org/projects/pyrddlgym/badge/?version=latest)](https://pyrddlgym.readthedocs.io/en/latest/jax.html)
+![License: MIT](https://img.shields.io/badge/License-MIT-blue.svg)
+[![Cumulative PyPI Downloads](https://img.shields.io/pypi/dm/pyrddlgym-jax)](https://pypistats.org/packages/pyrddlgym-jax)
+
+[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.** 
+
+Purpose:
 
 1. automatic translation of any RDDL description file into a differentiable simulator in JAX
 2. flexible policy class representations, automatic model relaxations for working in discrete and hybrid domains, and Bayesian hyper-parameter tuning.
@@ -25,17 +35,6 @@ 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).
-
-## Contents
-
-- [Installation](#installation)
-- [Running from the Command Line](#running-from-the-command-line)
-- [Running from Another Python Application](#running-from-another-python-application)
-- [Configuring the Planner](#configuring-the-planner)
-- [JaxPlan Dashboard](#jaxplan-dashboard)
-- [Tuning the Planner](#tuning-the-planner)
-- [Simulation](#simulation)
-- [Citing JaxPlan](#citing-jaxplan)
   
 ## Installation
 

pyrddlgym_jax-1.1/pyRDDLGym_jax/__init__.py

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

{pyrddlgym_jax-1.0 → pyrddlgym_jax-1.1}/pyRDDLGym_jax/core/compiler.py

@@ -51,65 +51,65 @@ class JaxRDDLCompiler:
             return func
         return exact_func
         
-    EXACT_RDDL_TO_JAX_NEGATIVE = wrap_logic(ExactLogic.exact_unary_function(jnp.negative))
+    EXACT_RDDL_TO_JAX_NEGATIVE = wrap_logic.__func__(ExactLogic.exact_unary_function(jnp.negative))
     EXACT_RDDL_TO_JAX_ARITHMETIC = {
-        '+': wrap_logic(ExactLogic.exact_binary_function(jnp.add)),
-        '-': wrap_logic(ExactLogic.exact_binary_function(jnp.subtract)),
-        '*': wrap_logic(ExactLogic.exact_binary_function(jnp.multiply)),
-        '/': wrap_logic(ExactLogic.exact_binary_function(jnp.divide))
+        '+': wrap_logic.__func__(ExactLogic.exact_binary_function(jnp.add)),
+        '-': wrap_logic.__func__(ExactLogic.exact_binary_function(jnp.subtract)),
+        '*': wrap_logic.__func__(ExactLogic.exact_binary_function(jnp.multiply)),
+        '/': wrap_logic.__func__(ExactLogic.exact_binary_function(jnp.divide))
     }            
     
     EXACT_RDDL_TO_JAX_RELATIONAL = {
-        '>=': wrap_logic(ExactLogic.exact_binary_function(jnp.greater_equal)),
-        '<=': wrap_logic(ExactLogic.exact_binary_function(jnp.less_equal)),
-        '<': wrap_logic(ExactLogic.exact_binary_function(jnp.less)),
-        '>': wrap_logic(ExactLogic.exact_binary_function(jnp.greater)),
-        '==': wrap_logic(ExactLogic.exact_binary_function(jnp.equal)),
-        '~=': wrap_logic(ExactLogic.exact_binary_function(jnp.not_equal))
+        '>=': wrap_logic.__func__(ExactLogic.exact_binary_function(jnp.greater_equal)),
+        '<=': wrap_logic.__func__(ExactLogic.exact_binary_function(jnp.less_equal)),
+        '<': wrap_logic.__func__(ExactLogic.exact_binary_function(jnp.less)),
+        '>': wrap_logic.__func__(ExactLogic.exact_binary_function(jnp.greater)),
+        '==': wrap_logic.__func__(ExactLogic.exact_binary_function(jnp.equal)),
+        '~=': wrap_logic.__func__(ExactLogic.exact_binary_function(jnp.not_equal))
     }        
     
-    EXACT_RDDL_TO_JAX_LOGICAL_NOT = wrap_logic(ExactLogic.exact_unary_function(jnp.logical_not))
+    EXACT_RDDL_TO_JAX_LOGICAL_NOT = wrap_logic.__func__(ExactLogic.exact_unary_function(jnp.logical_not))
     EXACT_RDDL_TO_JAX_LOGICAL = {
-        '^': wrap_logic(ExactLogic.exact_binary_function(jnp.logical_and)),
-        '&': wrap_logic(ExactLogic.exact_binary_function(jnp.logical_and)),
-        '|': wrap_logic(ExactLogic.exact_binary_function(jnp.logical_or)),
-        '~': wrap_logic(ExactLogic.exact_binary_function(jnp.logical_xor)),
-        '=>': wrap_logic(ExactLogic.exact_binary_implies),
-        '<=>': wrap_logic(ExactLogic.exact_binary_function(jnp.equal))
+        '^': wrap_logic.__func__(ExactLogic.exact_binary_function(jnp.logical_and)),
+        '&': wrap_logic.__func__(ExactLogic.exact_binary_function(jnp.logical_and)),
+        '|': wrap_logic.__func__(ExactLogic.exact_binary_function(jnp.logical_or)),
+        '~': wrap_logic.__func__(ExactLogic.exact_binary_function(jnp.logical_xor)),
+        '=>': wrap_logic.__func__(ExactLogic.exact_binary_implies),
+        '<=>': wrap_logic.__func__(ExactLogic.exact_binary_function(jnp.equal))
     }    
     
     EXACT_RDDL_TO_JAX_AGGREGATION = {
-        'sum': wrap_logic(ExactLogic.exact_aggregation(jnp.sum)),
-        'avg': wrap_logic(ExactLogic.exact_aggregation(jnp.mean)),
-        'prod': wrap_logic(ExactLogic.exact_aggregation(jnp.prod)),
-        'minimum': wrap_logic(ExactLogic.exact_aggregation(jnp.min)),
-        'maximum': wrap_logic(ExactLogic.exact_aggregation(jnp.max)),
-        'forall': wrap_logic(ExactLogic.exact_aggregation(jnp.all)),
-        'exists': wrap_logic(ExactLogic.exact_aggregation(jnp.any)),
-        'argmin': wrap_logic(ExactLogic.exact_aggregation(jnp.argmin)),
-        'argmax': wrap_logic(ExactLogic.exact_aggregation(jnp.argmax))
+        'sum': wrap_logic.__func__(ExactLogic.exact_aggregation(jnp.sum)),
+        'avg': wrap_logic.__func__(ExactLogic.exact_aggregation(jnp.mean)),
+        'prod': wrap_logic.__func__(ExactLogic.exact_aggregation(jnp.prod)),
+        'minimum': wrap_logic.__func__(ExactLogic.exact_aggregation(jnp.min)),
+        'maximum': wrap_logic.__func__(ExactLogic.exact_aggregation(jnp.max)),
+        'forall': wrap_logic.__func__(ExactLogic.exact_aggregation(jnp.all)),
+        'exists': wrap_logic.__func__(ExactLogic.exact_aggregation(jnp.any)),
+        'argmin': wrap_logic.__func__(ExactLogic.exact_aggregation(jnp.argmin)),
+        'argmax': wrap_logic.__func__(ExactLogic.exact_aggregation(jnp.argmax))
     }
     
     EXACT_RDDL_TO_JAX_UNARY = {        
-        'abs': wrap_logic(ExactLogic.exact_unary_function(jnp.abs)),
-        'sgn': wrap_logic(ExactLogic.exact_unary_function(jnp.sign)),
-        'round': wrap_logic(ExactLogic.exact_unary_function(jnp.round)),
-        'floor': wrap_logic(ExactLogic.exact_unary_function(jnp.floor)),
-        'ceil': wrap_logic(ExactLogic.exact_unary_function(jnp.ceil)),
-        'cos': wrap_logic(ExactLogic.exact_unary_function(jnp.cos)),
-        'sin': wrap_logic(ExactLogic.exact_unary_function(jnp.sin)),
-        'tan': wrap_logic(ExactLogic.exact_unary_function(jnp.tan)),
-        'acos': wrap_logic(ExactLogic.exact_unary_function(jnp.arccos)),
-        'asin': wrap_logic(ExactLogic.exact_unary_function(jnp.arcsin)),
-        'atan': wrap_logic(ExactLogic.exact_unary_function(jnp.arctan)),
-        'cosh': wrap_logic(ExactLogic.exact_unary_function(jnp.cosh)),
-        'sinh': wrap_logic(ExactLogic.exact_unary_function(jnp.sinh)),
-        'tanh': wrap_logic(ExactLogic.exact_unary_function(jnp.tanh)),
-        'exp': wrap_logic(ExactLogic.exact_unary_function(jnp.exp)),
-        'ln': wrap_logic(ExactLogic.exact_unary_function(jnp.log)),
-        'sqrt': wrap_logic(ExactLogic.exact_unary_function(jnp.sqrt)),
-        'lngamma': wrap_logic(ExactLogic.exact_unary_function(scipy.special.gammaln)),
-        'gamma': wrap_logic(ExactLogic.exact_unary_function(scipy.special.gamma))
+        'abs': wrap_logic.__func__(ExactLogic.exact_unary_function(jnp.abs)),
+        'sgn': wrap_logic.__func__(ExactLogic.exact_unary_function(jnp.sign)),
+        'round': wrap_logic.__func__(ExactLogic.exact_unary_function(jnp.round)),
+        'floor': wrap_logic.__func__(ExactLogic.exact_unary_function(jnp.floor)),
+        'ceil': wrap_logic.__func__(ExactLogic.exact_unary_function(jnp.ceil)),
+        'cos': wrap_logic.__func__(ExactLogic.exact_unary_function(jnp.cos)),
+        'sin': wrap_logic.__func__(ExactLogic.exact_unary_function(jnp.sin)),
+        'tan': wrap_logic.__func__(ExactLogic.exact_unary_function(jnp.tan)),
+        'acos': wrap_logic.__func__(ExactLogic.exact_unary_function(jnp.arccos)),
+        'asin': wrap_logic.__func__(ExactLogic.exact_unary_function(jnp.arcsin)),
+        'atan': wrap_logic.__func__(ExactLogic.exact_unary_function(jnp.arctan)),
+        'cosh': wrap_logic.__func__(ExactLogic.exact_unary_function(jnp.cosh)),
+        'sinh': wrap_logic.__func__(ExactLogic.exact_unary_function(jnp.sinh)),
+        'tanh': wrap_logic.__func__(ExactLogic.exact_unary_function(jnp.tanh)),
+        'exp': wrap_logic.__func__(ExactLogic.exact_unary_function(jnp.exp)),
+        'ln': wrap_logic.__func__(ExactLogic.exact_unary_function(jnp.log)),
+        'sqrt': wrap_logic.__func__(ExactLogic.exact_unary_function(jnp.sqrt)),
+        'lngamma': wrap_logic.__func__(ExactLogic.exact_unary_function(scipy.special.gammaln)),
+        'gamma': wrap_logic.__func__(ExactLogic.exact_unary_function(scipy.special.gamma))
     }              
     
     @staticmethod
@@ -117,23 +117,23 @@ class JaxRDDLCompiler:
         return jnp.log(x) / jnp.log(y), params
     
     EXACT_RDDL_TO_JAX_BINARY = {
-        'div': wrap_logic(ExactLogic.exact_binary_function(jnp.floor_divide)),
-        'mod': wrap_logic(ExactLogic.exact_binary_function(jnp.mod)),
-        'fmod': wrap_logic(ExactLogic.exact_binary_function(jnp.mod)),
-        'min': wrap_logic(ExactLogic.exact_binary_function(jnp.minimum)),
-        'max': wrap_logic(ExactLogic.exact_binary_function(jnp.maximum)),
-        'pow': wrap_logic(ExactLogic.exact_binary_function(jnp.power)),
-        'log': wrap_logic(_jax_wrapped_calc_log_exact),
-        'hypot': wrap_logic(ExactLogic.exact_binary_function(jnp.hypot)),
+        'div': wrap_logic.__func__(ExactLogic.exact_binary_function(jnp.floor_divide)),
+        'mod': wrap_logic.__func__(ExactLogic.exact_binary_function(jnp.mod)),
+        'fmod': wrap_logic.__func__(ExactLogic.exact_binary_function(jnp.mod)),
+        'min': wrap_logic.__func__(ExactLogic.exact_binary_function(jnp.minimum)),
+        'max': wrap_logic.__func__(ExactLogic.exact_binary_function(jnp.maximum)),
+        'pow': wrap_logic.__func__(ExactLogic.exact_binary_function(jnp.power)),
+        'log': wrap_logic.__func__(_jax_wrapped_calc_log_exact.__func__),
+        'hypot': wrap_logic.__func__(ExactLogic.exact_binary_function(jnp.hypot)),
     }
     
-    EXACT_RDDL_TO_JAX_IF = wrap_logic(ExactLogic.exact_if_then_else)
-    EXACT_RDDL_TO_JAX_SWITCH = wrap_logic(ExactLogic.exact_switch)
+    EXACT_RDDL_TO_JAX_IF = wrap_logic.__func__(ExactLogic.exact_if_then_else)
+    EXACT_RDDL_TO_JAX_SWITCH = wrap_logic.__func__(ExactLogic.exact_switch)
     
-    EXACT_RDDL_TO_JAX_BERNOULLI = wrap_logic(ExactLogic.exact_bernoulli)
-    EXACT_RDDL_TO_JAX_DISCRETE = wrap_logic(ExactLogic.exact_discrete)
-    EXACT_RDDL_TO_JAX_POISSON = wrap_logic(ExactLogic.exact_poisson)
-    EXACT_RDDL_TO_JAX_GEOMETRIC = wrap_logic(ExactLogic.exact_geometric)
+    EXACT_RDDL_TO_JAX_BERNOULLI = wrap_logic.__func__(ExactLogic.exact_bernoulli)
+    EXACT_RDDL_TO_JAX_DISCRETE = wrap_logic.__func__(ExactLogic.exact_discrete)
+    EXACT_RDDL_TO_JAX_POISSON = wrap_logic.__func__(ExactLogic.exact_poisson)
+    EXACT_RDDL_TO_JAX_GEOMETRIC = wrap_logic.__func__(ExactLogic.exact_geometric)
 
     def __init__(self, rddl: RDDLLiftedModel,
                  allow_synchronous_state: bool=True,

{pyrddlgym_jax-1.0 → pyrddlgym_jax-1.1}/pyRDDLGym_jax/core/planner.py

@@ -65,9 +65,8 @@ def _parse_config_file(path: str):
     config = configparser.RawConfigParser()
     config.optionxform = str 
     config.read(path)
-    args = {k: literal_eval(v) 
-            for section in config.sections()
-            for (k, v) in config.items(section)}
+    args = {section: {k: literal_eval(v) for (k, v) in config.items(section)}
+            for section in config.sections()}
     return config, args
 
 
@@ -75,9 +74,8 @@ def _parse_config_string(value: str):
     config = configparser.RawConfigParser()
     config.optionxform = str 
     config.read_string(value)
-    args = {k: literal_eval(v) 
-            for section in config.sections()
-            for (k, v) in config.items(section)}
+    args = {section: {k: literal_eval(v) for (k, v) in config.items(section)}
+            for section in config.sections()}
     return config, args
 
 
@@ -90,9 +88,9 @@ def _getattr_any(packages, item):
 
 
 def _load_config(config, args):
-    model_args = {k: args[k] for (k, _) in config.items('Model')}
-    planner_args = {k: args[k] for (k, _) in config.items('Optimizer')}
-    train_args = {k: args[k] for (k, _) in config.items('Training')}    
+    model_args = {k: args['Model'][k] for (k, _) in config.items('Model')}
+    planner_args = {k: args['Optimizer'][k] for (k, _) in config.items('Optimizer')}
+    train_args = {k: args['Training'][k] for (k, _) in config.items('Training')}    
     
     # read the model settings
     logic_name = model_args.get('logic', 'FuzzyLogic')
@@ -1661,7 +1659,7 @@ r"""
     def optimize_generator(self, key: Optional[random.PRNGKey]=None,
                            epochs: int=999999,
                            train_seconds: float=120.,
-                           dashboard: Optional[JaxPlannerDashboard]=None,
+                           dashboard: Optional[Any]=None,
                            dashboard_id: Optional[str]=None,
                            model_params: Optional[Dict[str, Any]]=None,
                            policy_hyperparams: Optional[Dict[str, Any]]=None,

{pyrddlgym_jax-1.0 → pyrddlgym_jax-1.1}/pyRDDLGym_jax/core/tuning.py

@@ -4,6 +4,7 @@ import threading
 import multiprocessing
 import os
 import time
+import traceback
 from typing import Any, Callable, Dict, Iterable, Optional, Tuple
 import warnings
 warnings.filterwarnings("ignore")
@@ -14,6 +15,7 @@ from bayes_opt.acquisition import AcquisitionFunction, UpperConfidenceBound
 import jax
 import numpy as np
 
+from pyRDDLGym.core.debug.exception import raise_warning
 from pyRDDLGym.core.env import RDDLEnv
 
 from pyRDDLGym_jax.core.planner import (
@@ -64,6 +66,7 @@ class JaxParameterTuning:
                  hyperparams: Hyperparameters,
                  online: bool,
                  eval_trials: int=5,
+                 rollouts_per_trial: int=1,
                  verbose: bool=True,
                  timeout_tuning: float=np.inf,
                  pool_context: str='spawn',
@@ -87,6 +90,8 @@ class JaxParameterTuning:
         hyperparameters in general (in seconds)
         :param eval_trials: how many trials to perform independent training
         in order to estimate the return for each set of hyper-parameters
+        :param rollouts_per_trial: how many rollouts to perform during evaluation
+        at the end of each training trial (only applies when online=False)
         :param verbose: whether to print intermediate results of tuning
         :param pool_context: context for multiprocessing pool (default "spawn")
         :param num_workers: how many points to evaluate in parallel
@@ -108,6 +113,7 @@ class JaxParameterTuning:
         self.hyperparams_dict = hyperparams_dict
         self.online = online
         self.eval_trials = eval_trials
+        self.rollouts_per_trial = rollouts_per_trial
         self.verbose = verbose
         
         # Bayesian parameters
@@ -154,6 +160,7 @@ class JaxParameterTuning:
               f'    mp_pool_poll_frequency    ={self.poll_frequency}\n'
               f'meta-objective parameters:\n'
               f'    planning_trials_per_iter  ={self.eval_trials}\n'
+              f'    rollouts_per_trial        ={self.rollouts_per_trial}\n'
               f'    acquisition_fn            ={self.acquisition}')
         
     @staticmethod
@@ -200,12 +207,14 @@ class JaxParameterTuning:
     
     @staticmethod
     def offline_trials(env, planner, train_args, key, iteration, index, num_trials, 
-                       verbose, viz, queue):
+                       rollouts_per_trial, verbose, viz, queue):
         average_reward = 0.0
         for trial in range(num_trials):
             key, subkey = jax.random.split(key)
+            
+            # for the dashboard
             experiment_id = f'iter={iteration}, worker={index}, trial={trial}'
-            if queue is not None:
+            if queue is not None and JaxPlannerDashboard is not None:
                 queue.put((
                     experiment_id, 
                     JaxPlannerDashboard.get_planner_info(planner), 
@@ -224,7 +233,8 @@ class JaxParameterTuning:
             policy = JaxOfflineController(
                 planner=planner, key=subkey, tqdm_position=index, 
                 params=best_params, train_on_reset=False)
-            total_reward = policy.evaluate(env, seed=np.array(subkey)[0])['mean']
+            total_reward = policy.evaluate(env, episodes=rollouts_per_trial, 
+                                           seed=np.array(subkey)[0])['mean']
             
             # update average reward
             if verbose:
@@ -243,8 +253,10 @@ class JaxParameterTuning:
         average_reward = 0.0
         for trial in range(num_trials):
             key, subkey = jax.random.split(key)
+            
+            # for the dashboard
             experiment_id = f'iter={iteration}, worker={index}, trial={trial}'
-            if queue is not None:
+            if queue is not None and JaxPlannerDashboard is not None:
                 queue.put((
                     experiment_id, 
                     JaxPlannerDashboard.get_planner_info(planner), 
@@ -304,6 +316,7 @@ class JaxParameterTuning:
         domain = kwargs['domain']
         instance = kwargs['instance']
         num_trials = kwargs['eval_trials']
+        rollouts_per_trial = kwargs['rollouts_per_trial']
         viz = kwargs['viz']
         verbose = kwargs['verbose']
         
@@ -332,7 +345,7 @@ class JaxParameterTuning:
         else:
             average_reward = JaxParameterTuning.offline_trials(
                 env, planner, train_args, key, iteration, index, 
-                num_trials, verbose, viz, queue
+                num_trials, rollouts_per_trial, verbose, viz, queue
             )
         
         pid = os.getpid()
@@ -353,7 +366,7 @@ class JaxParameterTuning:
             writer.writerow(COLUMNS + list(self.hyperparams_dict.keys()))
         
         # create a dash-board for visualizing experiment runs
-        if show_dashboard:
+        if show_dashboard and JaxPlannerDashboard is not None:
             dashboard = JaxPlannerDashboard()
             dashboard.launch()
             
@@ -365,6 +378,7 @@ class JaxParameterTuning:
             'domain': self.env.domain_text,
             'instance': self.env.instance_text,
             'eval_trials': self.eval_trials,
+            'rollouts_per_trial': self.rollouts_per_trial,
             'viz': self.env._visualizer,
             'verbose': self.verbose
         }

{pyrddlgym_jax-1.0 → pyrddlgym_jax-1.1}/pyRDDLGym_jax/core/visualization.py

@@ -1405,7 +1405,7 @@ class JaxPlannerDashboard:
         self.test_reward_dist[experiment_id] = callback['reward']
         self.train_state_fluents[experiment_id] = {
             name: np.asarray(callback['train_log']['fluents'][name])
-            for name in rddl.state_fluents or name in rddl.observ_fluents
+            for name in rddl.state_fluents
         }
         self.test_state_fluents[experiment_id] = {
             name: np.asarray(callback['fluents'][name])

{pyrddlgym_jax-1.0 → pyrddlgym_jax-1.1}/pyRDDLGym_jax.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pyRDDLGym-jax
-Version: 1.0
+Version: 1.1
 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
@@ -31,7 +31,17 @@ Requires-Dist: dash-bootstrap-components>=1.6.0; extra == "dashboard"
 
 # pyRDDLGym-jax
 
-**pyRDDLGym-jax (known in the literature as JaxPlan) is an efficient gradient-based/differentiable planning algorithm in JAX.** It provides:
+![Python Version](https://img.shields.io/badge/python-3.9%2B-blue)
+[![PyPI Version](https://img.shields.io/pypi/v/pyRDDLGym-jax.svg)](https://pypi.org/project/pyRDDLGym-jax/)
+[![Documentation Status](https://readthedocs.org/projects/pyrddlgym/badge/?version=latest)](https://pyrddlgym.readthedocs.io/en/latest/jax.html)
+![License: MIT](https://img.shields.io/badge/License-MIT-blue.svg)
+[![Cumulative PyPI Downloads](https://img.shields.io/pypi/dm/pyrddlgym-jax)](https://pypistats.org/packages/pyrddlgym-jax)
+
+[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.** 
+
+Purpose:
 
 1. automatic translation of any RDDL description file into a differentiable simulator in JAX
 2. flexible policy class representations, automatic model relaxations for working in discrete and hybrid domains, and Bayesian hyper-parameter tuning.
@@ -56,17 +66,6 @@ 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).
-
-## Contents
-
-- [Installation](#installation)
-- [Running from the Command Line](#running-from-the-command-line)
-- [Running from Another Python Application](#running-from-another-python-application)
-- [Configuring the Planner](#configuring-the-planner)
-- [JaxPlan Dashboard](#jaxplan-dashboard)
-- [Tuning the Planner](#tuning-the-planner)
-- [Simulation](#simulation)
-- [Citing JaxPlan](#citing-jaxplan)
   
 ## Installation
 

{pyrddlgym_jax-1.0 → pyrddlgym_jax-1.1}/setup.py

@@ -19,7 +19,7 @@ long_description = (Path(__file__).parent / "README.md").read_text()
 
 setup(
       name='pyRDDLGym-jax',
-      version='1.0',
+      version='1.1',
       author="Michael Gimelfarb, Ayal Taitler, Scott Sanner",
       author_email="mike.gimelfarb@mail.utoronto.ca, ataitler@gmail.com, ssanner@mie.utoronto.ca",
       description="pyRDDLGym-jax: automatic differentiation for solving sequential planning problems in JAX.",

pyrddlgym_jax-1.0/pyRDDLGym_jax/__init__.py

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