pyRDDLGym-jax 2.0__py3-none-any.whl → 2.1__py3-none-any.whl

pyRDDLGym_jax/core/planner.py

@@ -69,8 +69,7 @@ try:
     from pyRDDLGym_jax.core.visualization import JaxPlannerDashboard
 except Exception:
     raise_warning('Failed to load the dashboard visualization tool: '
-                  'please make sure you have installed the required packages.', 
-                  'red')
+                  'please make sure you have installed the required packages.', 'red')
     traceback.print_exc()
     JaxPlannerDashboard = None
 
@@ -133,7 +132,7 @@ def _load_config(config, args):
         comp_kwargs = model_args.get('complement_kwargs', {})
         compare_name = model_args.get('comparison', 'SigmoidComparison')
         compare_kwargs = model_args.get('comparison_kwargs', {})
-        sampling_name = model_args.get('sampling', 'GumbelSoftmax')
+        sampling_name = model_args.get('sampling', 'SoftRandomSampling')
         sampling_kwargs = model_args.get('sampling_kwargs', {})
         rounding_name = model_args.get('rounding', 'SoftRounding')
         rounding_kwargs = model_args.get('rounding_kwargs', {})
@@ -156,8 +155,7 @@ def _load_config(config, args):
         initializer = _getattr_any(
             packages=[initializers, hk.initializers], item=plan_initializer)
         if initializer is None:
-            raise_warning(
-                f'Ignoring invalid initializer <{plan_initializer}>.', 'red')
+            raise_warning(f'Ignoring invalid initializer <{plan_initializer}>.', 'red')
             del plan_kwargs['initializer']
         else:
             init_kwargs = plan_kwargs.pop('initializer_kwargs', {})
@@ -174,8 +172,7 @@ def _load_config(config, args):
         activation = _getattr_any(
             packages=[jax.nn, jax.numpy], item=plan_activation)
         if activation is None:
-            raise_warning(
-                f'Ignoring invalid activation <{plan_activation}>.', 'red')
+            raise_warning(f'Ignoring invalid activation <{plan_activation}>.', 'red')
             del plan_kwargs['activation']
         else:
             plan_kwargs['activation'] = activation
@@ -189,8 +186,7 @@ def _load_config(config, args):
     if planner_optimizer is not None:
         optimizer = _getattr_any(packages=[optax], item=planner_optimizer)
         if optimizer is None:
-            raise_warning(
-                f'Ignoring invalid optimizer <{planner_optimizer}>.', 'red')
+            raise_warning(f'Ignoring invalid optimizer <{planner_optimizer}>.', 'red')
             del planner_args['optimizer']
         else:
             planner_args['optimizer'] = optimizer
@@ -285,48 +281,14 @@ class JaxRDDLCompilerWithGrad(JaxRDDLCompiler):
         pvars_cast = set()
         for (var, values) in self.init_values.items():
             self.init_values[var] = np.asarray(values, dtype=self.REAL) 
-            if not np.issubdtype(np.atleast_1d(values).dtype, np.floating):
+            if not np.issubdtype(np.result_type(values), np.floating):
                 pvars_cast.add(var)
         if pvars_cast:
             raise_warning(f'JAX gradient compiler requires that initial values '
                           f'of p-variables {pvars_cast} be cast to float.')   
         
         # overwrite basic operations with fuzzy ones
-        self.RELATIONAL_OPS = {
-            '>=': logic.greater_equal,
-            '<=': logic.less_equal,
-            '<': logic.less,
-            '>': logic.greater,
-            '==': logic.equal,
-            '~=': logic.not_equal
-        }
-        self.LOGICAL_NOT = logic.logical_not
-        self.LOGICAL_OPS = {
-            '^': logic.logical_and,
-            '&': logic.logical_and,
-            '|': logic.logical_or,
-            '~': logic.xor,
-            '=>': logic.implies,
-            '<=>': logic.equiv
-        }
-        self.AGGREGATION_OPS['forall'] = logic.forall
-        self.AGGREGATION_OPS['exists'] = logic.exists
-        self.AGGREGATION_OPS['argmin'] = logic.argmin
-        self.AGGREGATION_OPS['argmax'] = logic.argmax
-        self.KNOWN_UNARY['sgn'] = logic.sgn
-        self.KNOWN_UNARY['floor'] = logic.floor
-        self.KNOWN_UNARY['ceil'] = logic.ceil
-        self.KNOWN_UNARY['round'] = logic.round
-        self.KNOWN_UNARY['sqrt'] = logic.sqrt
-        self.KNOWN_BINARY['div'] = logic.div
-        self.KNOWN_BINARY['mod'] = logic.mod
-        self.KNOWN_BINARY['fmod'] = logic.mod
-        self.IF_HELPER = logic.control_if
-        self.SWITCH_HELPER = logic.control_switch
-        self.BERNOULLI_HELPER = logic.bernoulli
-        self.DISCRETE_HELPER = logic.discrete
-        self.POISSON_HELPER = logic.poisson
-        self.GEOMETRIC_HELPER = logic.geometric
+        self.OPS = logic.get_operator_dicts()
         
     def _jax_stop_grad(self, jax_expr):        
         def _jax_wrapped_stop_grad(x, params, key):
@@ -575,7 +537,7 @@ class JaxStraightLinePlan(JaxPlan):
         def _jax_non_bool_param_to_action(var, param, hyperparams):
             if wrap_non_bool:
                 lower, upper = bounds_safe[var]
-                mb, ml, mu, mn = [mask.astype(compiled.REAL) 
+                mb, ml, mu, mn = [jnp.asarray(mask, dtype=compiled.REAL) 
                                   for mask in cond_lists[var]]       
                 action = (
                     mb * (lower + (upper - lower) * jax.nn.sigmoid(param)) + 
@@ -660,7 +622,7 @@ class JaxStraightLinePlan(JaxPlan):
                     action = _jax_non_bool_param_to_action(var, action, hyperparams)
                     action = jnp.clip(action, *bounds[var])
                     if ranges[var] == 'int':
-                        action = jnp.round(action).astype(compiled.INT)
+                        action = jnp.asarray(jnp.round(action), dtype=compiled.INT)
                     actions[var] = action
             return actions
         
@@ -961,12 +923,11 @@ class JaxDeepReactivePolicy(JaxPlan):
             non_bool_dims = 0
             for (var, values) in observed_vars.items():
                 if ranges[var] != 'bool':
-                    value_size = np.atleast_1d(values).size
+                    value_size = np.size(values)
                     if normalize_per_layer and value_size == 1:
                         raise_warning(
                             f'Cannot apply layer norm to state-fluent <{var}> '
-                            f'of size 1: setting normalize_per_layer = False.',
-                            'red')
+                            f'of size 1: setting normalize_per_layer = False.', 'red')
                         normalize_per_layer = False
                     non_bool_dims += value_size
             if not normalize_per_layer and non_bool_dims == 1:
@@ -990,9 +951,11 @@ class JaxDeepReactivePolicy(JaxPlan):
                     else:
                         if normalize and normalize_per_layer:
                             normalizer = hk.LayerNorm(
-                                axis=-1, param_axis=-1,
+                                axis=-1, 
+                                param_axis=-1,
                                 name=f'input_norm_{input_names[var]}',
-                                **self._normalizer_kwargs)
+                                **self._normalizer_kwargs
+                            )
                             state = normalizer(state)
                         states_non_bool.append(state)
                         non_bool_dims += state.size
@@ -1001,8 +964,11 @@ class JaxDeepReactivePolicy(JaxPlan):
             # optionally perform layer normalization on the non-bool inputs
             if normalize and not normalize_per_layer and non_bool_dims:
                 normalizer = hk.LayerNorm(
-                    axis=-1, param_axis=-1, name='input_norm',
-                    **self._normalizer_kwargs)
+                    axis=-1, 
+                    param_axis=-1, 
+                    name='input_norm',
+                    **self._normalizer_kwargs
+                )
                 normalized = normalizer(state[:non_bool_dims])
                 state = state.at[:non_bool_dims].set(normalized)
             return state
@@ -1021,7 +987,8 @@ class JaxDeepReactivePolicy(JaxPlan):
             actions = {}
             for (var, size) in layer_sizes.items():
                 linear = hk.Linear(size, name=layer_names[var], w_init=init)
-                reshape = hk.Reshape(output_shape=shapes[var], preserve_dims=-1,
+                reshape = hk.Reshape(output_shape=shapes[var], 
+                                     preserve_dims=-1,
                                      name=f'reshape_{layer_names[var]}')
                 output = reshape(linear(hidden))
                 if not shapes[var]:
@@ -1034,7 +1001,7 @@ class JaxDeepReactivePolicy(JaxPlan):
                 else:
                     if wrap_non_bool:
                         lower, upper = bounds_safe[var]
-                        mb, ml, mu, mn = [mask.astype(compiled.REAL) 
+                        mb, ml, mu, mn = [jnp.asarray(mask, dtype=compiled.REAL) 
                                           for mask in cond_lists[var]]       
                         action = (
                             mb * (lower + (upper - lower) * jax.nn.sigmoid(output)) + 
@@ -1048,8 +1015,7 @@ class JaxDeepReactivePolicy(JaxPlan):
             
             # for constraint satisfaction wrap bool actions with softmax
             if use_constraint_satisfaction:
-                linear = hk.Linear(
-                    bool_action_count, name='output_bool', w_init=init)
+                linear = hk.Linear(bool_action_count, name='output_bool', w_init=init)
                 output = jax.nn.softmax(linear(hidden))
                 actions[bool_key] = output
              
@@ -1087,8 +1053,7 @@ class JaxDeepReactivePolicy(JaxPlan):
         
         # test action prediction
         def _jax_wrapped_drp_predict_test(key, params, hyperparams, step, subs):
-            actions = _jax_wrapped_drp_predict_train(
-                key, params, hyperparams, step, subs)
+            actions = _jax_wrapped_drp_predict_train(key, params, hyperparams, step, subs)
             new_actions = {}
             for (var, action) in actions.items():
                 prange = ranges[var]
@@ -1096,7 +1061,7 @@ class JaxDeepReactivePolicy(JaxPlan):
                     new_action = action > 0.5
                 elif prange == 'int':
                     action = jnp.clip(action, *bounds[var])
-                    new_action = jnp.round(action).astype(compiled.INT)
+                    new_action = jnp.asarray(jnp.round(action), dtype=compiled.INT)
                 else:
                     new_action = jnp.clip(action, *bounds[var])
                 new_actions[var] = new_action
@@ -1436,8 +1401,8 @@ class GaussianPGPE(PGPE):
                     _jax_wrapped_pgpe_grad, 
                     in_axes=(0, None, None, None, None, None, None)
                 )(keys, mu, sigma, r_max, policy_hyperparams, subs, model_params)
-                mu_grad = jax.tree_map(partial(jnp.mean, axis=0), mu_grads)
-                sigma_grad = jax.tree_map(partial(jnp.mean, axis=0), sigma_grads)
+                mu_grad, sigma_grad = jax.tree_map(
+                    partial(jnp.mean, axis=0), (mu_grads, sigma_grads))
                 new_r_max = jnp.max(r_maxs)
             return mu_grad, sigma_grad, new_r_max
 
@@ -1463,6 +1428,71 @@ class GaussianPGPE(PGPE):
         self._update = jax.jit(_jax_wrapped_pgpe_update)
 
 
+# ***********************************************************************
+# ALL VERSIONS OF RISK FUNCTIONS
+# 
+# Based on the original paper "A Distributional Framework for Risk-Sensitive 
+# End-to-End Planning in Continuous MDPs" by Patton et al., AAAI 2022.
+#
+# Original risk functions:
+# - entropic utility
+# - mean-variance
+# - mean-semideviation
+# - conditional value at risk with straight-through gradient trick
+#
+# ***********************************************************************
+
+
+@jax.jit
+def entropic_utility(returns: jnp.ndarray, beta: float) -> float:
+    return (-1.0 / beta) * jax.scipy.special.logsumexp(
+        -beta * returns, b=1.0 / returns.size)
+
+
+@jax.jit
+def mean_variance_utility(returns: jnp.ndarray, beta: float) -> float:
+    return jnp.mean(returns) - 0.5 * beta * jnp.var(returns)
+
+
+@jax.jit
+def mean_deviation_utility(returns: jnp.ndarray, beta: float) -> float:
+    return jnp.mean(returns) - 0.5 * beta * jnp.std(returns)
+
+
+@jax.jit
+def mean_semideviation_utility(returns: jnp.ndarray, beta: float) -> float:
+    mu = jnp.mean(returns)
+    msd = jnp.sqrt(jnp.mean(jnp.minimum(0.0, returns - mu) ** 2))
+    return mu - 0.5 * beta * msd
+
+
+@jax.jit
+def mean_semivariance_utility(returns: jnp.ndarray, beta: float) -> float:
+    mu = jnp.mean(returns)
+    msv = jnp.mean(jnp.minimum(0.0, returns - mu) ** 2)
+    return mu - 0.5 * beta * msv
+
+
+@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)
+
+
+UTILITY_LOOKUP = {
+    'mean': jnp.mean,
+    'mean_var': mean_variance_utility,
+    'mean_std': mean_deviation_utility,
+    'mean_semivar': mean_semivariance_utility,
+    'mean_semidev': mean_semideviation_utility,
+    'entropic': entropic_utility,
+    'exponential': entropic_utility,
+    'cvar': cvar_utility
+}
+
+
 # ***********************************************************************
 # ALL VERSIONS OF JAX PLANNER
 # 
@@ -1525,8 +1555,7 @@ class JaxBackpropPlanner:
         reward as a form of normalization
         :param utility: how to aggregate return observations to compute utility
         of a policy or plan; must be either a function mapping jax array to a 
-        scalar, or a a string identifying the utility function by name 
-        ("mean", "mean_var", "entropic", or "cvar" are currently supported)
+        scalar, or a a string identifying the utility function by name
         :param utility_kwargs: additional keyword arguments to pass hyper-
         parameters to the utility function call
         :param cpfs_without_grad: which CPFs do not have gradients (use straight
@@ -1584,18 +1613,11 @@ class JaxBackpropPlanner:
         # set utility
         if isinstance(utility, str):
             utility = utility.lower()
-            if utility == 'mean':
-                utility_fn = jnp.mean
-            elif utility == 'mean_var':
-                utility_fn = mean_variance_utility
-            elif utility == 'entropic':
-                utility_fn = entropic_utility
-            elif utility == 'cvar':
-                utility_fn = cvar_utility
-            else:
+            utility_fn = UTILITY_LOOKUP.get(utility, None)
+            if utility_fn is None:
                 raise RDDLNotImplementedError(
-                    f'Utility function <{utility}> is not supported: '
-                    'must be one of ["mean", "mean_var", "entropic", "cvar"].')
+                    f'Utility <{utility}> is not supported, '
+                    f'must be one of {list(UTILITY_LOOKUP.keys())}.')
         else:
             utility_fn = utility
         self.utility = utility_fn
@@ -1865,7 +1887,7 @@ r"""
                     f'{set(self.test_compiled.init_values.keys())}.')
             value = np.reshape(value, newshape=np.shape(init_value))[np.newaxis, ...]
             train_value = np.repeat(value, repeats=n_train, axis=0)
-            train_value = train_value.astype(self.compiled.REAL)
+            train_value = np.asarray(train_value, dtype=self.compiled.REAL)
             init_train[name] = train_value
             init_test[name] = np.repeat(value, repeats=n_test, axis=0)
         
@@ -2175,7 +2197,9 @@ r"""
         
         iters = range(epochs)
         if print_progress:
-            iters = tqdm(iters, total=100, position=tqdm_position)
+            iters = tqdm(iters, total=100, 
+                         bar_format='{l_bar}{bar}| {elapsed} {postfix}', 
+                         position=tqdm_position)
         position_str = '' if tqdm_position is None else f'[{tqdm_position}]'
         
         for it in iters:
@@ -2256,7 +2280,8 @@ r"""
                 status = JaxPlannerStatus.ITER_BUDGET_REACHED
             
             # build a callback
-            progress_percent = int(100 * min(1, max(elapsed / train_seconds, it / epochs)))
+            progress_percent = 100 * min(
+                1, max(0, elapsed / train_seconds, it / (epochs - 1)))
             callback = {
                 'status': status,
                 'iteration': it,
@@ -2279,7 +2304,7 @@ r"""
                 'train_log': train_log,
                 **test_log
             }
-            
+
             # stopping condition reached
             if stopping_rule is not None and stopping_rule.monitor(callback):
                 callback['status'] = status = JaxPlannerStatus.STOPPING_RULE_REACHED  
@@ -2290,8 +2315,10 @@ r"""
                 iters.set_description(
                     f'{position_str} {it:6} it / {-train_loss:14.5f} train / '
                     f'{-test_loss_smooth:14.5f} test / {-best_loss:14.5f} best / '
-                    f'{status.value} status / {total_pgpe_it:6} pgpe'
+                    f'{status.value} status / {total_pgpe_it:6} pgpe',
+                    refresh=False
                 )
+                iters.set_postfix_str(f"{(it + 1) / elapsed:.2f}it/s", refresh=True)
             
             # dash-board
             if dashboard is not None:
@@ -2332,7 +2359,7 @@ r"""
                 last_iter_improve, -train_loss, -test_loss_smooth, -best_loss, grad_norm)
             print(f'summary of optimization:\n'
                   f'    status        ={status}\n'
-                  f'    time          ={elapsed:.6f} sec.\n'
+                  f'    time          ={elapsed:.3f} sec.\n'
                   f'    iterations    ={it}\n'
                   f'    best objective={-best_loss:.6f}\n'
                   f'    best grad norm={grad_norm}\n'
@@ -2358,12 +2385,12 @@ r"""
                 return termcolor.colored(
                     '[FAILURE] no progress was made '
                     f'and max grad norm {max_grad_norm:.6f} was zero: '
-                    'the solver was likely stuck in a plateau.', 'red')
+                    'solver likely stuck in a plateau.', 'red')
             else:
                 return termcolor.colored(
                     '[FAILURE] no progress was made '
                     f'but max grad norm {max_grad_norm:.6f} was non-zero: '
-                    'the learning rate or other hyper-parameters were likely suboptimal.', 
+                    'learning rate or other hyper-parameters likely suboptimal.', 
                     'red')
         
         # model is likely poor IF:
@@ -2372,8 +2399,8 @@ r"""
             return termcolor.colored(
                 '[WARNING] progress was made '
                 f'but relative train-test error {validation_error:.6f} was high: '
-                'model relaxation around the solution was poor '
-                'or the batch size was too small.', 'yellow')
+                'poor model relaxation around solution or batch size too small.', 
+                'yellow')
         
         # model likely did not converge IF:
         # 1. the max grad relative to the return is high
@@ -2383,9 +2410,9 @@ r"""
                 return termcolor.colored(
                     '[WARNING] progress was made '
                     f'but max grad norm {max_grad_norm:.6f} was high: '
-                    'the solution was likely locally suboptimal, '
-                    'or the relaxed model was not smooth around the solution, '
-                    'or the batch size was too small.', 'yellow')
+                    'solution locally suboptimal '
+                    'or relaxed model not smooth around solution '
+                    'or batch size too small.', 'yellow')
         
         # likely successful
         return termcolor.colored(
@@ -2412,8 +2439,7 @@ r"""
         for (var, values) in subs.items():
             
             # must not be grounded
-            if RDDLPlanningModel.FLUENT_SEP in var \
-            or RDDLPlanningModel.OBJECT_SEP in var:
+            if RDDLPlanningModel.FLUENT_SEP in var or RDDLPlanningModel.OBJECT_SEP in var:
                 raise ValueError(f'State dictionary passed to the JAX policy is '
                                  f'grounded, since it contains the key <{var}>, '
                                  f'but a vectorized environment is required: '
@@ -2421,9 +2447,8 @@ r"""
             
             # must be numeric array
             # exception is for POMDPs at 1st epoch when observ-fluents are None
-            dtype = np.atleast_1d(values).dtype
-            if not np.issubdtype(dtype, np.number) \
-            and not np.issubdtype(dtype, np.bool_):
+            dtype = np.result_type(values)
+            if not np.issubdtype(dtype, np.number) and not np.issubdtype(dtype, np.bool_):
                 if step == 0 and var in self.rddl.observ_fluents:
                     subs[var] = self.test_compiled.init_values[var]
                 else:
@@ -2435,40 +2460,7 @@ r"""
         actions = self.test_policy(key, params, policy_hyperparams, step, subs)
         actions = jax.tree_map(np.asarray, actions)
         return actions      
-    
-
-# ***********************************************************************
-# ALL VERSIONS OF RISK FUNCTIONS
-# 
-# Based on the original paper "A Distributional Framework for Risk-Sensitive 
-# End-to-End Planning in Continuous MDPs" by Patton et al., AAAI 2022.
-#
-# Original risk functions:
-# - entropic utility
-# - mean-variance approximation
-# - conditional value at risk with straight-through gradient trick
-#
-# ***********************************************************************
-
-
-@jax.jit
-def entropic_utility(returns: jnp.ndarray, beta: float) -> float:
-    return (-1.0 / beta) * jax.scipy.special.logsumexp(
-        -beta * returns, b=1.0 / returns.size)
-
-
-@jax.jit
-def mean_variance_utility(returns: jnp.ndarray, beta: float) -> float:
-    return jnp.mean(returns) - 0.5 * beta * jnp.var(returns)
-    
-
-@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)
-   
+       
 
 # ***********************************************************************
 # ALL VERSIONS OF CONTROLLERS
@@ -2580,8 +2572,7 @@ class JaxOnlineController(BaseAgent):
         self.callback = callback
         params = callback['best_params']
         self.key, subkey = random.split(self.key)
-        actions = planner.get_action(
-            subkey, params, 0, state, self.eval_hyperparams)
+        actions = planner.get_action(subkey, params, 0, state, self.eval_hyperparams)
         if self.warm_start:
             self.guess = planner.plan.guess_next_epoch(params)
         return actions

pyRDDLGym_jax/core/visualization.py

@@ -20,8 +20,7 @@ import math
 import numpy as np
 import time
 import threading
-from typing import Any, Dict, List, Optional, Tuple, TYPE_CHECKING
-import warnings
+from typing import Any, Dict, Optional, Tuple, TYPE_CHECKING
 import webbrowser
 
 # prevent endless console prints
@@ -32,7 +31,7 @@ log.setLevel(logging.ERROR)
 import dash
 from dash.dcc import Interval, Graph, Store
 from dash.dependencies import Input, Output, State, ALL
-from dash.html import Div, B, H4, P, Img, Hr
+from dash.html import Div, B, H4, P, Hr
 import dash_bootstrap_components as dbc
 
 import plotly.colors as pc 
@@ -53,6 +52,7 @@ REWARD_ERROR_DIST_SUBPLOTS = 20
 MODEL_STATE_ERROR_HEIGHT = 300
 POLICY_STATE_VIZ_MAX_HEIGHT = 800
 GP_POSTERIOR_MAX_HEIGHT = 800
+GP_POSTERIOR_PIXELS = 100
 
 PLOT_AXES_FONT_SIZE = 11
 EXPERIMENT_ENTRY_FONT_SIZE = 14
@@ -1417,7 +1417,7 @@ class JaxPlannerDashboard:
         self.pgpe_return[experiment_id].append(callback['pgpe_return'])
         
         # data for return distributions
-        progress = callback['progress']
+        progress = int(callback['progress'])
         if progress - self.return_dist_last_progress[experiment_id] \
             >= PROGRESS_FOR_NEXT_RETURN_DIST:
             self.return_dist_ticks[experiment_id].append(iteration)
@@ -1486,8 +1486,8 @@ class JaxPlannerDashboard:
                 if i2 > i1:
                     
                     # Generate a grid for visualization
-                    p1_values = np.linspace(*bounds[param1], 100)
-                    p2_values = np.linspace(*bounds[param2], 100)
+                    p1_values = np.linspace(*bounds[param1], GP_POSTERIOR_PIXELS)
+                    p2_values = np.linspace(*bounds[param2], GP_POSTERIOR_PIXELS)
                     P1, P2 = np.meshgrid(p1_values, p2_values)
                     
                     # Predict the mean and deviation of the surrogate model
@@ -1500,8 +1500,7 @@ class JaxPlannerDashboard:
                     for p1, p2 in zip(np.ravel(P1), np.ravel(P2)):
                         params = {param1: p1, param2: p2}
                         params.update(fixed_params)
-                        param_grid.append(
-                            [params[key] for key in optimizer.space.keys])
+                        param_grid.append([params[key] for key in optimizer.space.keys])
                     param_grid = np.asarray(param_grid)
                     mean, std = optimizer._gp.predict(param_grid, return_std=True)
                     mean = mean.reshape(P1.shape)

pyRDDLGym_jax/examples/run_tune.py

@@ -3,7 +3,7 @@ is performed using a batched parallelized Bayesian optimization.
 
 The syntax is:
 
-    python run_tune.py <domain> <instance> <method> [<trials>] [<iters>] [<workers>]
+    python run_tune.py <domain> <instance> <method> [<trials>] [<iters>] [<workers>] [<dashboard>]
     
 where:
     <domain> is the name of a domain located in the /Examples directory
@@ -15,6 +15,7 @@ where:
     (defaults to 20)
     <workers> is the number of parallel workers (i.e. batch size), which must
     not exceed the number of cores available on the machine (defaults to 4)
+    <dashboard> is whether the dashboard is displayed
 '''
 import os
 import sys
@@ -35,7 +36,7 @@ def power_10(x):
     return 10.0 ** x
 
 
-def main(domain, instance, method, trials=5, iters=20, workers=4):
+def main(domain, instance, method, trials=5, iters=20, workers=4, dashboard=False):
     
     # set up the environment   
     env = pyRDDLGym.make(domain, instance, vectorized=True)
@@ -48,9 +49,9 @@ def main(domain, instance, method, trials=5, iters=20, workers=4):
     
     # map parameters in the config that will be tuned
     hyperparams = [
-        Hyperparameter('MODEL_WEIGHT_TUNE', -1., 5., power_10),
+        Hyperparameter('MODEL_WEIGHT_TUNE', -1., 4., power_10),
         Hyperparameter('POLICY_WEIGHT_TUNE', -2., 2., power_10),
-        Hyperparameter('LEARNING_RATE_TUNE', -5., 1., power_10),
+        Hyperparameter('LEARNING_RATE_TUNE', -5., 0., power_10),
         Hyperparameter('LAYER1_TUNE', 1, 8, power_2),
         Hyperparameter('LAYER2_TUNE', 1, 8, power_2),
         Hyperparameter('ROLLOUT_HORIZON_TUNE', 1, min(env.horizon, 100), int)       
@@ -64,7 +65,9 @@ def main(domain, instance, method, trials=5, iters=20, workers=4):
                                 eval_trials=trials,
                                 num_workers=workers,
                                 gp_iters=iters)
-    tuning.tune(key=42, log_file=f'gp_{method}_{domain}_{instance}.csv')
+    tuning.tune(key=42, 
+                log_file=f'gp_{method}_{domain}_{instance}.csv', 
+                show_dashboard=dashboard)
     
     # evaluate the agent on the best parameters
     planner_args, _, train_args = load_config_from_string(tuning.best_config)
@@ -77,7 +80,7 @@ def main(domain, instance, method, trials=5, iters=20, workers=4):
 
 def run_from_args(args):
     if len(args) < 3:
-        print('python run_tune.py <domain> <instance> <method> [<trials>] [<iters>] [<workers>]')
+        print('python run_tune.py <domain> <instance> <method> [<trials>] [<iters>] [<workers>] [<dashboard>]')
         exit(1)
     if args[2] not in ['drp', 'slp', 'replan']:
         print('<method> in [drp, slp, replan]')
@@ -86,6 +89,7 @@ def run_from_args(args):
     if len(args) >= 4: kwargs['trials'] = int(args[3])
     if len(args) >= 5: kwargs['iters'] = int(args[4])
     if len(args) >= 6: kwargs['workers'] = int(args[5])
+    if len(args) >= 7: kwargs['dashboard'] = bool(args[6])
     main(**kwargs)