pyRDDLGym-jax 0.4__py3-none-any.whl → 1.0__py3-none-any.whl

pyRDDLGym_jax/core/planner.py

@@ -31,17 +31,18 @@ from pyRDDLGym.core.policy import BaseAgent
 from pyRDDLGym_jax import __version__
 from pyRDDLGym_jax.core import logic
 from pyRDDLGym_jax.core.compiler import JaxRDDLCompiler
-from pyRDDLGym_jax.core.logic import FuzzyLogic
+from pyRDDLGym_jax.core.logic import Logic, FuzzyLogic
 
-# try to import matplotlib, if failed then skip plotting
+# try to load the dash board
 try:
-    import matplotlib.pyplot as plt
+    from pyRDDLGym_jax.core.visualization import JaxPlannerDashboard
 except Exception:
-    raise_warning('failed to import matplotlib: '
-                  'plotting functionality will be disabled.', 'red')
+    raise_warning('Failed to load the dashboard visualization tool: '
+                  'please make sure you have installed the required packages.', 
+                  'red')
     traceback.print_exc()
-    plt = None
-    
+    JaxPlannerDashboard = None
+
 Activation = Callable[[jnp.ndarray], jnp.ndarray]
 Bounds = Dict[str, Tuple[np.ndarray, np.ndarray]]
 Kwargs = Dict[str, Any]
@@ -57,6 +58,7 @@ Pytree = Any
 #
 # ***********************************************************************
 
+
 def _parse_config_file(path: str):
     if not os.path.isfile(path):
         raise FileNotFoundError(f'File {path} does not exist.')
@@ -95,18 +97,25 @@ def _load_config(config, args):
     # read the model settings
     logic_name = model_args.get('logic', 'FuzzyLogic')
     logic_kwargs = model_args.get('logic_kwargs', {})
-    tnorm_name = model_args.get('tnorm', 'ProductTNorm')
-    tnorm_kwargs = model_args.get('tnorm_kwargs', {})
-    comp_name = model_args.get('complement', 'StandardComplement')
-    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_kwargs = model_args.get('sampling_kwargs', {})
-    logic_kwargs['tnorm'] = getattr(logic, tnorm_name)(**tnorm_kwargs)
-    logic_kwargs['complement'] = getattr(logic, comp_name)(**comp_kwargs)
-    logic_kwargs['comparison'] = getattr(logic, compare_name)(**compare_kwargs)
-    logic_kwargs['sampling'] = getattr(logic, sampling_name)(**sampling_kwargs)
+    if logic_name == 'FuzzyLogic':
+        tnorm_name = model_args.get('tnorm', 'ProductTNorm')
+        tnorm_kwargs = model_args.get('tnorm_kwargs', {})
+        comp_name = model_args.get('complement', 'StandardComplement')
+        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_kwargs = model_args.get('sampling_kwargs', {})
+        rounding_name = model_args.get('rounding', 'SoftRounding')
+        rounding_kwargs = model_args.get('rounding_kwargs', {})
+        control_name = model_args.get('control', 'SoftControlFlow')
+        control_kwargs = model_args.get('control_kwargs', {})
+        logic_kwargs['tnorm'] = getattr(logic, tnorm_name)(**tnorm_kwargs)
+        logic_kwargs['complement'] = getattr(logic, comp_name)(**comp_kwargs)
+        logic_kwargs['comparison'] = getattr(logic, compare_name)(**compare_kwargs)
+        logic_kwargs['sampling'] = getattr(logic, sampling_name)(**sampling_kwargs)
+        logic_kwargs['rounding'] = getattr(logic, rounding_name)(**rounding_kwargs)
+        logic_kwargs['control'] = getattr(logic, control_name)(**control_kwargs)
     
     # read the policy settings
     plan_method = planner_args.pop('method')
@@ -157,11 +166,25 @@ def _load_config(config, args):
         else:
             planner_args['optimizer'] = optimizer
         
-    # read the optimize call settings
+    # optimize call RNG key
     planner_key = train_args.get('key', None)
     if planner_key is not None:
         train_args['key'] = random.PRNGKey(planner_key)
     
+    # dashboard
+    dashboard_key = train_args.get('dashboard', None)
+    if dashboard_key is not None and dashboard_key and JaxPlannerDashboard is not None:
+        train_args['dashboard'] = JaxPlannerDashboard()
+    elif dashboard_key is not None:
+        del train_args['dashboard']        
+    
+    # optimize call stopping rule
+    stopping_rule = train_args.get('stopping_rule', None)
+    if stopping_rule is not None:
+        stopping_rule_kwargs = train_args.pop('stopping_rule_kwargs', {})
+        train_args['stopping_rule'] = getattr(
+            sys.modules[__name__], stopping_rule)(**stopping_rule_kwargs)
+    
     return planner_args, plan_kwargs, train_args
 
 
@@ -175,7 +198,7 @@ def load_config_from_string(value: str) -> Tuple[Kwargs, ...]:
     '''Loads config file contents specified explicitly as a string value.'''
     config, args = _parse_config_string(value)
     return _load_config(config, args)
-
+    
     
 # ***********************************************************************
 # MODEL RELAXATIONS
@@ -193,7 +216,7 @@ class JaxRDDLCompilerWithGrad(JaxRDDLCompiler):
     '''
     
     def __init__(self, *args,
-                 logic: FuzzyLogic=FuzzyLogic(),
+                 logic: Logic=FuzzyLogic(),
                  cpfs_without_grad: Optional[Set[str]]=None,
                  **kwargs) -> None:
         '''Creates a new RDDL to Jax compiler, where operations that are not
@@ -228,57 +251,55 @@ class JaxRDDLCompilerWithGrad(JaxRDDLCompiler):
         
         # overwrite basic operations with fuzzy ones
         self.RELATIONAL_OPS = {
-            '>=': logic.greater_equal(),
-            '<=': logic.less_equal(),
-            '<': logic.less(),
-            '>': logic.greater(),
-            '==': logic.equal(),
-            '~=': logic.not_equal()
+            '>=': logic.greater_equal,
+            '<=': logic.less_equal,
+            '<': logic.less,
+            '>': logic.greater,
+            '==': logic.equal,
+            '~=': logic.not_equal
         }
-        self.LOGICAL_NOT = logic.logical_not()
+        self.LOGICAL_NOT = logic.logical_not
         self.LOGICAL_OPS = {
-            '^': logic.logical_and(),
-            '&': logic.logical_and(),
-            '|': logic.logical_or(),
-            '~': logic.xor(),
-            '=>': logic.implies(),
-            '<=>': logic.equiv()
+            '^': 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()
-        
-    def _jax_stop_grad(self, jax_expr):
-        
+        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
+        
+    def _jax_stop_grad(self, jax_expr):        
         def _jax_wrapped_stop_grad(x, params, key):
-            sample, key, error = jax_expr(x, params, key)
+            sample, key, error, params = jax_expr(x, params, key)
             sample = jax.lax.stop_gradient(sample)
-            return sample, key, error
-        
+            return sample, key, error, params
         return _jax_wrapped_stop_grad
         
-    def _compile_cpfs(self, info):
+    def _compile_cpfs(self, init_params):
         cpfs_cast = set()   
         jax_cpfs = {}
         for (_, cpfs) in self.levels.items():
             for cpf in cpfs:
                 _, expr = self.rddl.cpfs[cpf]
-                jax_cpfs[cpf] = self._jax(expr, info, dtype=self.REAL)
+                jax_cpfs[cpf] = self._jax(expr, init_params, dtype=self.REAL)
                 if self.rddl.variable_ranges[cpf] != 'real':
                     cpfs_cast.add(cpf)
                 if cpf in self.cpfs_without_grad:
@@ -289,17 +310,15 @@ class JaxRDDLCompilerWithGrad(JaxRDDLCompiler):
                           f'{cpfs_cast} be cast to float.') 
         if self.cpfs_without_grad:
             raise_warning(f'User requested that gradients not flow '
-                          f'through CPFs {self.cpfs_without_grad}.')      
+                          f'through CPFs {self.cpfs_without_grad}.')    
+ 
         return jax_cpfs
     
-    def _jax_kron(self, expr, info):
-        if self.logic.verbose:
-            raise_warning('JAX gradient compiler ignores KronDelta '
-                          'during compilation.')            
+    def _jax_kron(self, expr, init_params):       
         arg, = expr.args
-        arg = self._jax(arg, info)
+        arg = self._jax(arg, init_params)
         return arg
-    
+
 
 # ***********************************************************************
 # ALL VERSIONS OF JAX PLANS
@@ -309,6 +328,7 @@ class JaxRDDLCompilerWithGrad(JaxRDDLCompiler):
 #
 # ***********************************************************************
 
+
 class JaxPlan:
     '''Base class for all JAX policy representations.'''
     
@@ -320,7 +340,7 @@ class JaxPlan:
         self.bounds = None
         
     def summarize_hyperparameters(self) -> None:
-        pass
+        print(self.__str__())
         
     def compile(self, compiled: JaxRDDLCompilerWithGrad,
                 _bounds: Bounds,
@@ -363,7 +383,7 @@ class JaxPlan:
         self._projection = value
     
     def _calculate_action_info(self, compiled: JaxRDDLCompilerWithGrad,
-                               user_bounds: Bounds, 
+                               user_bounds: Bounds,
                                horizon: int):
         shapes, bounds, bounds_safe, cond_lists = {}, {}, {}, {}
         for (name, prange) in compiled.rddl.variable_ranges.items():
@@ -446,24 +466,24 @@ class JaxStraightLinePlan(JaxPlan):
         self._wrap_softmax = wrap_softmax
         self._use_new_projection = use_new_projection
         self._max_constraint_iter = max_constraint_iter
-        
-    def summarize_hyperparameters(self) -> None:
+    
+    def __str__(self) -> str:
         bounds = '\n        '.join(
             map(lambda kv: f'{kv[0]}: {kv[1]}', self.bounds.items()))
-        print(f'policy hyper-parameters:\n'
-              f'    initializer          ={self._initializer_base}\n'
-              f'constraint-sat strategy (simple):\n'
-              f'    parsed_action_bounds =\n        {bounds}\n'
-              f'    wrap_sigmoid         ={self._wrap_sigmoid}\n'
-              f'    wrap_sigmoid_min_prob={self._min_action_prob}\n'
-              f'    wrap_non_bool        ={self._wrap_non_bool}\n'
-              f'constraint-sat strategy (complex):\n'
-              f'    wrap_softmax         ={self._wrap_softmax}\n'
-              f'    use_new_projection   ={self._use_new_projection}\n'
-              f'    max_projection_iters ={self._max_constraint_iter}')
+        return (f'policy hyper-parameters:\n'
+                f'    initializer          ={self._initializer_base}\n'
+                f'constraint-sat strategy (simple):\n'
+                f'    parsed_action_bounds =\n        {bounds}\n'
+                f'    wrap_sigmoid         ={self._wrap_sigmoid}\n'
+                f'    wrap_sigmoid_min_prob={self._min_action_prob}\n'
+                f'    wrap_non_bool        ={self._wrap_non_bool}\n'
+                f'constraint-sat strategy (complex):\n'
+                f'    wrap_softmax         ={self._wrap_softmax}\n'
+                f'    use_new_projection   ={self._use_new_projection}\n'
+                f'    max_projection_iters ={self._max_constraint_iter}')
     
     def compile(self, compiled: JaxRDDLCompilerWithGrad,
-                _bounds: Bounds, 
+                _bounds: Bounds,
                 horizon: int) -> None:
         rddl = compiled.rddl
         
@@ -504,7 +524,7 @@ class JaxStraightLinePlan(JaxPlan):
         def _jax_bool_action_to_param(var, action, hyperparams):
             if wrap_sigmoid:
                 weight = hyperparams[var]
-                return (-1.0 / weight) * jnp.log(1.0 / action - 1.0)
+                return jax.scipy.special.logit(action) / weight
             else:
                 return action
             
@@ -513,14 +533,13 @@ class JaxStraightLinePlan(JaxPlan):
         def _jax_non_bool_param_to_action(var, param, hyperparams):
             if wrap_non_bool:
                 lower, upper = bounds_safe[var]
-                action = jnp.select(
-                    condlist=cond_lists[var],
-                    choicelist=[
-                        lower + (upper - lower) * jax.nn.sigmoid(param),
-                        lower + (jax.nn.elu(param) + 1.0),
-                        upper - (jax.nn.elu(-param) + 1.0),
-                        param
-                    ]
+                mb, ml, mu, mn = [mask.astype(compiled.REAL) 
+                                  for mask in cond_lists[var]]       
+                action = (
+                    mb * (lower + (upper - lower) * jax.nn.sigmoid(param)) + 
+                    ml * (lower + (jax.nn.elu(param) + 1.0)) + 
+                    mu * (upper - (jax.nn.elu(-param) + 1.0)) + 
+                    mn * param
                 )
             else:
                 action = param
@@ -780,7 +799,7 @@ class JaxDeepReactivePolicy(JaxPlan):
     def __init__(self, topology: Optional[Sequence[int]]=None,
                  activation: Activation=jnp.tanh,
                  initializer: hk.initializers.Initializer=hk.initializers.VarianceScaling(scale=2.0),
-                 normalize: bool=False, 
+                 normalize: bool=False,
                  normalize_per_layer: bool=False,
                  normalizer_kwargs: Optional[Kwargs]=None,
                  wrap_non_bool: bool=False) -> None:
@@ -812,23 +831,23 @@ class JaxDeepReactivePolicy(JaxPlan):
             normalizer_kwargs = {'create_offset': True, 'create_scale': True}
         self._normalizer_kwargs = normalizer_kwargs
         self._wrap_non_bool = wrap_non_bool
-            
-    def summarize_hyperparameters(self) -> None:
+    
+    def __str__(self) -> str:
         bounds = '\n        '.join(
             map(lambda kv: f'{kv[0]}: {kv[1]}', self.bounds.items()))
-        print(f'policy hyper-parameters:\n'
-              f'    topology            ={self._topology}\n'
-              f'    activation_fn       ={self._activations[0].__name__}\n'
-              f'    initializer         ={type(self._initializer_base).__name__}\n'
-              f'    apply_input_norm    ={self._normalize}\n'
-              f'    input_norm_layerwise={self._normalize_per_layer}\n'
-              f'    input_norm_args     ={self._normalizer_kwargs}\n'
-              f'constraint-sat strategy:\n'
-              f'    parsed_action_bounds=\n        {bounds}\n'
-              f'    wrap_non_bool       ={self._wrap_non_bool}')
-    
+        return (f'policy hyper-parameters:\n'
+                f'    topology            ={self._topology}\n'
+                f'    activation_fn       ={self._activations[0].__name__}\n'
+                f'    initializer         ={type(self._initializer_base).__name__}\n'
+                f'    apply_input_norm    ={self._normalize}\n'
+                f'    input_norm_layerwise={self._normalize_per_layer}\n'
+                f'    input_norm_args     ={self._normalizer_kwargs}\n'
+                f'constraint-sat strategy:\n'
+                f'    parsed_action_bounds=\n        {bounds}\n'
+                f'    wrap_non_bool       ={self._wrap_non_bool}')
+        
     def compile(self, compiled: JaxRDDLCompilerWithGrad,
-                _bounds: Bounds, 
+                _bounds: Bounds,
                 horizon: int) -> None:
         rddl = compiled.rddl
         
@@ -881,7 +900,7 @@ class JaxDeepReactivePolicy(JaxPlan):
                     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.', 
+                            f'of size 1: setting normalize_per_layer = False.',
                             'red')
                         normalize_per_layer = False
                     non_bool_dims += value_size
@@ -906,8 +925,8 @@ class JaxDeepReactivePolicy(JaxPlan):
                     else:
                         if normalize and normalize_per_layer:
                             normalizer = hk.LayerNorm(
-                                axis=-1, param_axis=-1, 
-                                name=f'input_norm_{input_names[var]}', 
+                                axis=-1, param_axis=-1,
+                                name=f'input_norm_{input_names[var]}',
                                 **self._normalizer_kwargs)
                             state = normalizer(state)
                         states_non_bool.append(state)
@@ -917,7 +936,7 @@ 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', 
+                    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)
@@ -950,14 +969,13 @@ class JaxDeepReactivePolicy(JaxPlan):
                 else:
                     if wrap_non_bool:
                         lower, upper = bounds_safe[var]
-                        action = jnp.select(
-                            condlist=cond_lists[var],
-                            choicelist=[
-                                lower + (upper - lower) * jax.nn.sigmoid(output),
-                                lower + (jax.nn.elu(output) + 1.0),
-                                upper - (jax.nn.elu(-output) + 1.0),
-                                output
-                            ]
+                        mb, ml, mu, mn = [mask.astype(compiled.REAL) 
+                                          for mask in cond_lists[var]]       
+                        action = (
+                            mb * (lower + (upper - lower) * jax.nn.sigmoid(output)) + 
+                            ml * (lower + (jax.nn.elu(output) + 1.0)) + 
+                            mu * (upper - (jax.nn.elu(-output) + 1.0)) + 
+                            mn * output
                         )
                     else:
                         action = output
@@ -1049,7 +1067,7 @@ class JaxDeepReactivePolicy(JaxPlan):
         
     def guess_next_epoch(self, params: Pytree) -> Pytree:
         return params
-
+    
     
 # ***********************************************************************
 # ALL VERSIONS OF JAX PLANNER
@@ -1059,6 +1077,7 @@ class JaxDeepReactivePolicy(JaxPlan):
 #
 # ***********************************************************************
 
+
 class RollingMean:
     '''Maintains an estimate of the rolling mean of a stream of real-valued 
     observations.'''
@@ -1077,113 +1096,6 @@ class RollingMean:
         return self._total / len(memory)
 
 
-class JaxPlannerPlot:
-    '''Supports plotting and visualization of a JAX policy in real time.'''
-    
-    def __init__(self, rddl: RDDLPlanningModel, horizon: int, 
-                 show_violin: bool=True, show_action: bool=True) -> None:
-        '''Creates a new planner visualizer.
-        
-        :param rddl: the planning model to optimize
-        :param horizon: the lookahead or planning horizon
-        :param show_violin: whether to show the distribution of batch losses
-        :param show_action: whether to show heatmaps of the action fluents
-        '''
-        num_plots = 1
-        if show_violin:
-            num_plots += 1
-        if show_action:
-            num_plots += len(rddl.action_fluents)
-        self._fig, axes = plt.subplots(num_plots)
-        if num_plots == 1:
-            axes = [axes]
-        
-        # prepare the loss plot
-        self._loss_ax = axes[0]
-        self._loss_ax.autoscale(enable=True)
-        self._loss_ax.set_xlabel('training time')
-        self._loss_ax.set_ylabel('loss value')
-        self._loss_plot = self._loss_ax.plot(
-            [], [], linestyle=':', marker='o', markersize=2)[0]
-        self._loss_back = self._fig.canvas.copy_from_bbox(self._loss_ax.bbox)
-        
-        # prepare the violin plot
-        if show_violin:
-            self._hist_ax = axes[1]
-        else:
-            self._hist_ax = None
-            
-        # prepare the action plots
-        if show_action:
-            self._action_ax = {name: axes[idx + (2 if show_violin else 1)]
-                               for (idx, name) in enumerate(rddl.action_fluents)}
-            self._action_plots = {}
-            for name in rddl.action_fluents:
-                ax = self._action_ax[name]
-                if rddl.variable_ranges[name] == 'bool':
-                    vmin, vmax = 0.0, 1.0
-                else:
-                    vmin, vmax = None, None  
-                action_dim = 1
-                for dim in rddl.object_counts(rddl.variable_params[name]):
-                    action_dim *= dim     
-                action_plot = ax.pcolormesh(
-                    np.zeros((action_dim, horizon)), 
-                    cmap='seismic', vmin=vmin, vmax=vmax)
-                ax.set_aspect('auto')        
-                ax.set_xlabel('decision epoch')
-                ax.set_ylabel(name)
-                plt.colorbar(action_plot, ax=ax)
-                self._action_plots[name] = action_plot
-            self._action_back = {name: self._fig.canvas.copy_from_bbox(ax.bbox)
-                                 for (name, ax) in self._action_ax.items()}
-        else:
-            self._action_ax = None
-            self._action_plots = None
-            self._action_back = None
-            
-        plt.tight_layout()
-        plt.show(block=False)
-        
-    def redraw(self, xticks, losses, actions, returns) -> None:
-        
-        # draw the loss curve
-        self._fig.canvas.restore_region(self._loss_back)
-        self._loss_plot.set_xdata(xticks)
-        self._loss_plot.set_ydata(losses)        
-        self._loss_ax.set_xlim([0, len(xticks)])
-        self._loss_ax.set_ylim([np.min(losses), np.max(losses)])
-        self._loss_ax.draw_artist(self._loss_plot)
-        self._fig.canvas.blit(self._loss_ax.bbox)
-        
-        # draw the violin plot
-        if self._hist_ax is not None:
-            self._hist_ax.clear()
-            self._hist_ax.set_xlabel('loss value')
-            self._hist_ax.set_ylabel('density')
-            self._hist_ax.violinplot(returns, vert=False, showmeans=True)
-        
-        # draw the actions
-        if self._action_ax is not None:
-            for (name, values) in actions.items():
-                values = np.mean(values, axis=0, dtype=float)
-                values = np.reshape(values, newshape=(values.shape[0], -1)).T
-                self._fig.canvas.restore_region(self._action_back[name])
-                self._action_plots[name].set_array(values)
-                self._action_ax[name].draw_artist(self._action_plots[name])
-                self._fig.canvas.blit(self._action_ax[name].bbox)
-                self._action_plots[name].set_clim([np.min(values), np.max(values)])
-                
-        self._fig.canvas.draw()
-        self._fig.canvas.flush_events()
-        
-    def close(self) -> None:
-        plt.close(self._fig)
-        del self._loss_ax, self._hist_ax, self._action_ax, \
-            self._loss_plot, self._action_plots, self._fig, \
-            self._loss_back, self._action_back
-
-
 class JaxPlannerStatus(Enum):
     '''Represents the status of a policy update from the JAX planner, 
     including whether the update resulted in nan gradient, 
@@ -1191,16 +1103,50 @@ class JaxPlannerStatus(Enum):
     can be used to monitor and act based on the planner's progress.'''
     
     NORMAL = 0
-    NO_PROGRESS = 1
-    PRECONDITION_POSSIBLY_UNSATISFIED = 2
-    INVALID_GRADIENT = 3
-    TIME_BUDGET_REACHED = 4
-    ITER_BUDGET_REACHED = 5
+    STOPPING_RULE_REACHED = 1
+    NO_PROGRESS = 2
+    PRECONDITION_POSSIBLY_UNSATISFIED = 3
+    INVALID_GRADIENT = 4
+    TIME_BUDGET_REACHED = 5
+    ITER_BUDGET_REACHED = 6
     
-    def is_failure(self) -> bool:
-        return self.value >= 3
+    def is_terminal(self) -> bool:
+        return self.value == 1 or self.value >= 4
 
 
+class JaxPlannerStoppingRule:
+    '''The base class of all planner stopping rules.'''
+    
+    def reset(self) -> None:
+        raise NotImplementedError
+        
+    def monitor(self, callback: Dict[str, Any]) -> bool:
+        raise NotImplementedError
+    
+
+class NoImprovementStoppingRule(JaxPlannerStoppingRule):
+    '''Stopping rule based on no improvement for a fixed number of iterations.'''
+    
+    def __init__(self, patience: int) -> None:
+        self.patience = patience
+    
+    def reset(self) -> None:
+        self.callback = None
+        self.iters_since_last_update = 0
+        
+    def monitor(self, callback: Dict[str, Any]) -> bool:
+        if self.callback is None \
+        or callback['best_return'] > self.callback['best_return']:
+            self.callback = callback
+            self.iters_since_last_update = 0
+        else:
+            self.iters_since_last_update += 1
+        return self.iters_since_last_update >= self.patience
+    
+    def __str__(self) -> str:
+        return f'No improvement for {self.patience} iterations'
+        
+
 class JaxBackpropPlanner:
     '''A class for optimizing an action sequence in the given RDDL MDP using 
     gradient descent.'''
@@ -1215,13 +1161,16 @@ class JaxBackpropPlanner:
                  optimizer: Callable[..., optax.GradientTransformation]=optax.rmsprop,
                  optimizer_kwargs: Optional[Kwargs]=None,
                  clip_grad: Optional[float]=None,
-                 logic: FuzzyLogic=FuzzyLogic(),
+                 line_search_kwargs: Optional[Kwargs]=None,
+                 noise_kwargs: Optional[Kwargs]=None,
+                 logic: Logic=FuzzyLogic(),
                  use_symlog_reward: bool=False,
                  utility: Union[Callable[[jnp.ndarray], float], str]='mean',
                  utility_kwargs: Optional[Kwargs]=None,
                  cpfs_without_grad: Optional[Set[str]]=None,
                  compile_non_fluent_exact: bool=True,
-                 logger: Optional[Logger]=None) -> None:
+                 logger: Optional[Logger]=None,
+                 dashboard_viz: Optional[Any]=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
@@ -1241,7 +1190,10 @@ class JaxBackpropPlanner:
         :param optimizer_kwargs: a dictionary of parameters to pass to the SGD
         factory (e.g. which parameters are controllable externally)
         :param clip_grad: maximum magnitude of gradient updates
-        :param logic: a subclass of FuzzyLogic for mapping exact mathematical
+        :param line_search_kwargs: parameters to pass to optional line search
+        method to scale learning rate
+        :param noise_kwargs: parameters of optional gradient noise
+        :param logic: a subclass of Logic for mapping exact mathematical
         operations to their differentiable counterparts 
         :param use_symlog_reward: whether to use the symlog transform on the 
         reward as a form of normalization
@@ -1256,6 +1208,8 @@ class JaxBackpropPlanner:
         :param compile_non_fluent_exact: whether non-fluent expressions 
         are always compiled using exact JAX expressions
         :param logger: to log information about compilation to file
+        :param dashboard_viz: optional visualizer object from the environment
+        to pass to the dashboard to visualize the policy
         '''
         self.rddl = rddl
         self.plan = plan
@@ -1270,29 +1224,34 @@ class JaxBackpropPlanner:
             action_bounds = {}
         self._action_bounds = action_bounds
         self.use64bit = use64bit
-        self._optimizer_name = optimizer
+        self.optimizer_name = optimizer
         if optimizer_kwargs is None:
             optimizer_kwargs = {'learning_rate': 0.1}
-        self._optimizer_kwargs = optimizer_kwargs
+        self.optimizer_kwargs = optimizer_kwargs
         self.clip_grad = clip_grad
+        self.line_search_kwargs = line_search_kwargs
+        self.noise_kwargs = noise_kwargs
         
         # set optimizer
         try:
             optimizer = optax.inject_hyperparams(optimizer)(**optimizer_kwargs)
         except Exception as _:
             raise_warning(
-                'Failed to inject hyperparameters into optax optimizer, '
+                f'Failed to inject hyperparameters into optax optimizer {optimizer}, '
                 'rolling back to safer method: please note that modification of '
-                'optimizer hyperparameters will not work, and it is '
-                'recommended to update optax and related packages.', 'red')
-            optimizer = optimizer(**optimizer_kwargs)     
-        if clip_grad is None:
-            self.optimizer = optimizer
-        else:
-            self.optimizer = optax.chain(
-                optax.clip(clip_grad),
-                optimizer
-            )
+                'optimizer hyperparameters will not work.', 'red')
+            optimizer = optimizer(**optimizer_kwargs)   
+        
+        # apply optimizer chain of transformations
+        pipeline = []  
+        if clip_grad is not None:
+            pipeline.append(optax.clip(clip_grad))
+        if noise_kwargs is not None:
+            pipeline.append(optax.add_noise(**noise_kwargs))
+        pipeline.append(optimizer)
+        if line_search_kwargs is not None:
+            pipeline.append(optax.scale_by_zoom_linesearch(**line_search_kwargs))
+        self.optimizer = optax.chain(*pipeline)
         
         # set utility
         if isinstance(utility, str):
@@ -1325,11 +1284,12 @@ class JaxBackpropPlanner:
         self.cpfs_without_grad = cpfs_without_grad
         self.compile_non_fluent_exact = compile_non_fluent_exact
         self.logger = logger
+        self.dashboard_viz = dashboard_viz
         
         self._jax_compile_rddl()        
         self._jax_compile_optimizer()
     
-    def _summarize_system(self) -> None:
+    def summarize_system(self) -> str:
         try:
             jaxlib_version = jax._src.lib.version_str
         except Exception as _:
@@ -1340,42 +1300,63 @@ class JaxBackpropPlanner:
         except Exception as _:
             devices_short = 'N/A'
         LOGO = \
+r"""
+   __   ______   __  __   ______  __       ______   __   __    
+  /\ \ /\  __ \ /\_\_\_\ /\  == \/\ \     /\  __ \ /\ "-.\ \   
+ _\_\ \\ \  __ \\/_/\_\/_\ \  _-/\ \ \____\ \  __ \\ \ \-.  \  
+/\_____\\ \_\ \_\ /\_\/\_\\ \_\   \ \_____\\ \_\ \_\\ \_\\"\_\ 
+\/_____/ \/_/\/_/ \/_/\/_/ \/_/    \/_____/ \/_/\/_/ \/_/ \/_/ 
 """
-   __    ______    __  __    ______  __        ______    __   __    
-  /\ \  /\  __ \  /\_\_\_\  /\  == \/\ \      /\  __ \  /\ "-.\ \   
- _\_\ \ \ \  __ \ \/_/\_\/_ \ \  _-/\ \ \____ \ \  __ \ \ \ \-.  \  
-/\_____\ \ \_\ \_\  /\_\/\_\ \ \_\   \ \_____\ \ \_\ \_\ \ \_\\"\_\ 
-\/_____/  \/_/\/_/  \/_/\/_/  \/_/    \/_____/  \/_/\/_/  \/_/ \/_/ 
-"""
-                                                       
-        print('\n'
-              f'{LOGO}\n'
-              f'Version {__version__}\n' 
-              f'Python {sys.version}\n'
-              f'jax {jax.version.__version__}, jaxlib {jaxlib_version}, '
-              f'optax {optax.__version__}, haiku {hk.__version__}, '
-              f'numpy {np.__version__}\n'
-              f'devices: {devices_short}\n')
+                   
+        return ('\n'
+                f'{LOGO}\n'
+                f'Version {__version__}\n' 
+                f'Python {sys.version}\n'
+                f'jax {jax.version.__version__}, jaxlib {jaxlib_version}, '
+                f'optax {optax.__version__}, haiku {hk.__version__}, '
+                f'numpy {np.__version__}\n'
+                f'devices: {devices_short}\n')
+    
+    def __str__(self) -> str:
+        result = (f'objective hyper-parameters:\n'
+                  f'    utility_fn        ={self.utility.__name__}\n'
+                  f'    utility args      ={self.utility_kwargs}\n'
+                  f'    use_symlog        ={self.use_symlog_reward}\n'
+                  f'    lookahead         ={self.horizon}\n'
+                  f'    user_action_bounds={self._action_bounds}\n'
+                  f'    fuzzy logic type  ={type(self.logic).__name__}\n'
+                  f'    non_fluents exact ={self.compile_non_fluent_exact}\n'
+                  f'    cpfs_no_gradient  ={self.cpfs_without_grad}\n'
+                  f'optimizer hyper-parameters:\n'
+                  f'    use_64_bit        ={self.use64bit}\n'
+                  f'    optimizer         ={self.optimizer_name}\n'
+                  f'    optimizer args    ={self.optimizer_kwargs}\n'
+                  f'    clip_gradient     ={self.clip_grad}\n'
+                  f'    line_search_kwargs={self.line_search_kwargs}\n'
+                  f'    noise_kwargs      ={self.noise_kwargs}\n'
+                  f'    batch_size_train  ={self.batch_size_train}\n'
+                  f'    batch_size_test   ={self.batch_size_test}')
+        result += '\n' + str(self.plan)
+        result += '\n' + str(self.logic)
+        
+        # print model relaxation information
+        if not self.compiled.model_params:
+            return result
+        result += '\n' + ('Some RDDL operations are non-differentiable '
+                          'and will be approximated as follows:' + '\n')
+        exprs_by_rddl_op, values_by_rddl_op = {}, {}
+        for info in self.compiled.model_parameter_info().values():
+            rddl_op = info['rddl_op']
+            exprs_by_rddl_op.setdefault(rddl_op, []).append(info['id'])
+            values_by_rddl_op.setdefault(rddl_op, []).append(info['init_value'])
+        for rddl_op in sorted(exprs_by_rddl_op.keys()):
+            result += (f'    {rddl_op}:\n'
+                       f'        addresses  ={exprs_by_rddl_op[rddl_op]}\n'
+                       f'        init_values={values_by_rddl_op[rddl_op]}\n')
+        return result
         
     def summarize_hyperparameters(self) -> None:
-        print(f'objective hyper-parameters:\n'
-              f'    utility_fn        ={self.utility.__name__}\n'
-              f'    utility args      ={self.utility_kwargs}\n'
-              f'    use_symlog        ={self.use_symlog_reward}\n'
-              f'    lookahead         ={self.horizon}\n'
-              f'    user_action_bounds={self._action_bounds}\n'
-              f'    fuzzy logic type  ={type(self.logic).__name__}\n'
-              f'    nonfluents exact  ={self.compile_non_fluent_exact}\n'
-              f'    cpfs_no_gradient  ={self.cpfs_without_grad}\n'
-              f'optimizer hyper-parameters:\n'
-              f'    use_64_bit        ={self.use64bit}\n'
-              f'    optimizer         ={self._optimizer_name.__name__}\n'
-              f'    optimizer args    ={self._optimizer_kwargs}\n'
-              f'    clip_gradient     ={self.clip_grad}\n'
-              f'    batch_size_train  ={self.batch_size_train}\n'
-              f'    batch_size_test   ={self.batch_size_test}')
-        self.plan.summarize_hyperparameters()
-        self.logic.summarize_hyperparameters()
+        print(self.__str__())
         
     # ===========================================================================
     # COMPILATION SUBROUTINES
@@ -1391,14 +1372,16 @@ class JaxBackpropPlanner:
             logger=self.logger,
             use64bit=self.use64bit,
             cpfs_without_grad=self.cpfs_without_grad,
-            compile_non_fluent_exact=self.compile_non_fluent_exact)
+            compile_non_fluent_exact=self.compile_non_fluent_exact
+        )
         self.compiled.compile(log_jax_expr=True, heading='RELAXED MODEL')
         
         # Jax compilation of the exact RDDL for testing
         self.test_compiled = JaxRDDLCompiler(
-            rddl=rddl, 
+            rddl=rddl,
             logger=self.logger,
-            use64bit=self.use64bit)
+            use64bit=self.use64bit
+        )
         self.test_compiled.compile(log_jax_expr=True, heading='EXACT MODEL')
         
     def _jax_compile_optimizer(self):
@@ -1414,13 +1397,15 @@ class JaxBackpropPlanner:
         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
+        )
         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
+        )
         self.test_rollouts = jax.jit(test_rollouts)
         
         # initialization
@@ -1455,14 +1440,16 @@ class JaxBackpropPlanner:
         _jax_wrapped_returns = self._jax_return(use_symlog)
         
         # the loss is the average cumulative reward across all roll-outs
-        def _jax_wrapped_plan_loss(key, policy_params, hyperparams,
+        def _jax_wrapped_plan_loss(key, policy_params, policy_hyperparams,
                                    subs, model_params):
-            log = rollouts(key, policy_params, hyperparams, subs, model_params)
+            log, model_params = rollouts(
+                key, policy_params, policy_hyperparams, subs, model_params)
             rewards = log['reward']
             returns = _jax_wrapped_returns(rewards)
             utility = utility_fn(returns, **utility_kwargs)
             loss = -utility
-            return loss, log
+            aux = (log, model_params)
+            return loss, aux
         
         return _jax_wrapped_plan_loss
     
@@ -1470,30 +1457,45 @@ class JaxBackpropPlanner:
         init = self.plan.initializer
         optimizer = self.optimizer
         
-        def _jax_wrapped_init_policy(key, hyperparams, subs):
-            policy_params = init(key, hyperparams, subs)
+        # initialize both the policy and its optimizer
+        def _jax_wrapped_init_policy(key, policy_hyperparams, subs):
+            policy_params = init(key, policy_hyperparams, subs)
             opt_state = optimizer.init(policy_params)
-            return policy_params, opt_state, None
+            return policy_params, opt_state, {}
         
         return _jax_wrapped_init_policy
         
     def _jax_update(self, loss):
         optimizer = self.optimizer
         projection = self.plan.projection
+        use_ls = self.line_search_kwargs is not None
         
         # calculate the plan gradient w.r.t. return loss and update optimizer
         # also perform a projection step to satisfy constraints on actions
-        def _jax_wrapped_plan_update(key, policy_params, hyperparams,
+        def _jax_wrapped_loss_swapped(policy_params, key, policy_hyperparams,
+                                      subs, model_params):
+            return loss(key, policy_params, policy_hyperparams, subs, model_params)[0]
+            
+        def _jax_wrapped_plan_update(key, policy_params, policy_hyperparams,
                                      subs, model_params, opt_state, opt_aux):
             grad_fn = jax.value_and_grad(loss, argnums=1, has_aux=True)
-            (loss_val, log), grad = grad_fn(
-                key, policy_params, hyperparams, subs, model_params)  
-            updates, opt_state = optimizer.update(grad, opt_state) 
+            (loss_val, (log, model_params)), grad = grad_fn(
+                key, policy_params, policy_hyperparams, subs, model_params)
+            if use_ls:
+                updates, opt_state = optimizer.update(
+                    grad, opt_state, params=policy_params, 
+                    value=loss_val, grad=grad, value_fn=_jax_wrapped_loss_swapped,
+                    key=key, policy_hyperparams=policy_hyperparams, subs=subs,
+                    model_params=model_params)
+            else:
+                updates, opt_state = optimizer.update(
+                    grad, opt_state, params=policy_params) 
             policy_params = optax.apply_updates(policy_params, updates)
-            policy_params, converged = projection(policy_params, hyperparams)
+            policy_params, converged = projection(policy_params, policy_hyperparams)
             log['grad'] = grad
             log['updates'] = updates
-            return policy_params, converged, opt_state, None, loss_val, log
+            return policy_params, converged, opt_state, opt_aux, \
+                loss_val, log, model_params
         
         return jax.jit(_jax_wrapped_plan_update)
             
@@ -1520,11 +1522,10 @@ class JaxBackpropPlanner:
         for (state, next_state) in rddl.next_state.items():
             init_train[next_state] = init_train[state]
             init_test[next_state] = init_test[state]
-        
         return init_train, init_test
     
     def as_optimization_problem(
-            self, key: Optional[random.PRNGKey]=None, 
+            self, key: Optional[random.PRNGKey]=None,
             policy_hyperparams: Optional[Pytree]=None,
             loss_function_updates_key: bool=True,
             grad_function_updates_key: bool=False) -> Tuple[Callable, Callable, np.ndarray, Callable]:
@@ -1574,38 +1575,40 @@ class JaxBackpropPlanner:
         loss_fn = self._jax_loss(self.train_rollouts)
         
         @jax.jit
-        def _loss_with_key(key, params_1d):
+        def _loss_with_key(key, params_1d, model_params):
             policy_params = unravel_fn(params_1d)
-            loss_val, _ = loss_fn(key, policy_params, policy_hyperparams, 
-                                  train_subs, model_params)
-            return loss_val
+            loss_val, (_, model_params) = loss_fn(
+                key, policy_params, policy_hyperparams, train_subs, model_params)
+            return loss_val, model_params
         
         @jax.jit
-        def _grad_with_key(key, params_1d):
+        def _grad_with_key(key, params_1d, model_params):
             policy_params = unravel_fn(params_1d)
             grad_fn = jax.grad(loss_fn, argnums=1, has_aux=True)
-            grad_val, _ = grad_fn(key, policy_params, policy_hyperparams, 
-                                  train_subs, model_params)
-            grad_1d = jax.flatten_util.ravel_pytree(grad_val)[0]
-            return grad_1d
+            grad_val, (_, model_params) = grad_fn(
+                key, policy_params, policy_hyperparams, train_subs, model_params)
+            grad_val = jax.flatten_util.ravel_pytree(grad_val)[0]
+            return grad_val, model_params 
         
         def _loss_function(params_1d):
             nonlocal key
+            nonlocal model_params
             if loss_function_updates_key:
                 key, subkey = random.split(key)
             else:
                 subkey = key
-            loss_val = _loss_with_key(subkey, params_1d)
+            loss_val, model_params = _loss_with_key(subkey, params_1d, model_params)
             loss_val = float(loss_val)
             return loss_val
         
         def _grad_function(params_1d):
             nonlocal key
+            nonlocal model_params
             if grad_function_updates_key:
                 key, subkey = random.split(key)
             else:
                 subkey = key
-            grad = _grad_with_key(subkey, params_1d)
+            grad, model_params = _grad_with_key(subkey, params_1d, model_params)
             grad = np.asarray(grad)
             return grad
         
@@ -1620,9 +1623,9 @@ class JaxBackpropPlanner:
         
         :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 plot_step: frequency to plot the plan and save result to disk
-        :param plot_kwargs: additional arguments to pass to the plotter
+        :param train_seconds: total time allocated for gradient descent               
+        :param dashboard: dashboard to display training results
+        :param dashboard_id: experiment id for the dashboard
         :param model_params: optional model-parameters to override default
         :param policy_hyperparams: hyper-parameters for the policy/plan, such as
         weights for sigmoid wrapping boolean actions
@@ -1633,6 +1636,7 @@ class JaxBackpropPlanner:
         :param print_summary: whether to print planner header, parameter 
         summary, and diagnosis
         :param print_progress: whether to print the progress bar during training
+        :param stopping_rule: stopping criterion
         :param test_rolling_window: the test return is averaged on a rolling 
         window of the past test_rolling_window returns when updating the best
         parameters found so far
@@ -1657,14 +1661,15 @@ class JaxBackpropPlanner:
     def optimize_generator(self, key: Optional[random.PRNGKey]=None,
                            epochs: int=999999,
                            train_seconds: float=120.,
-                           plot_step: Optional[int]=None,
-                           plot_kwargs: Optional[Dict[str, Any]]=None,
+                           dashboard: Optional[JaxPlannerDashboard]=None,
+                           dashboard_id: Optional[str]=None,
                            model_params: Optional[Dict[str, Any]]=None,
                            policy_hyperparams: Optional[Dict[str, Any]]=None,
                            subs: Optional[Dict[str, Any]]=None,
                            guess: Optional[Pytree]=None,
                            print_summary: bool=True,
                            print_progress: bool=True,
+                           stopping_rule: Optional[JaxPlannerStoppingRule]=None,
                            test_rolling_window: int=10,
                            tqdm_position: Optional[int]=None) -> Generator[Dict[str, Any], None, None]:
         '''Returns a generator for computing an optimal policy or plan. 
@@ -1673,9 +1678,9 @@ class JaxBackpropPlanner:
         
         :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 plot_step: frequency to plot the plan and save result to disk
-        :param plot_kwargs: additional arguments to pass to the plotter
+        :param train_seconds: total time allocated for gradient descent        
+        :param dashboard: dashboard to display training results
+        :param dashboard_id: experiment id for the dashboard
         :param model_params: optional model-parameters to override default
         :param policy_hyperparams: hyper-parameters for the policy/plan, such as
         weights for sigmoid wrapping boolean actions
@@ -1686,6 +1691,7 @@ class JaxBackpropPlanner:
         :param print_summary: whether to print planner header, parameter 
         summary, and diagnosis
         :param print_progress: whether to print the progress bar during training
+        :param stopping_rule: stopping criterion
         :param test_rolling_window: the test return is averaged on a rolling 
         window of the past test_rolling_window returns when updating the best
         parameters found so far
@@ -1694,9 +1700,14 @@ class JaxBackpropPlanner:
         start_time = time.time()
         elapsed_outside_loop = 0
         
+        # ======================================================================
+        # INITIALIZATION OF HYPER-PARAMETERS
+        # ======================================================================
+
         # if PRNG key is not provided
         if key is None:
             key = random.PRNGKey(round(time.time() * 1000))
+        dash_key = key[1].item()
             
         # if policy_hyperparams is not provided
         if policy_hyperparams is None:
@@ -1723,7 +1734,7 @@ class JaxBackpropPlanner:
             
         # print summary of parameters:
         if print_summary:
-            self._summarize_system()
+            print(self.summarize_system())
             self.summarize_hyperparameters()
             print(f'optimize() call hyper-parameters:\n'
                   f'    PRNG key           ={key}\n'
@@ -1734,15 +1745,16 @@ class JaxBackpropPlanner:
                   f'    override_subs_dict ={subs is not None}\n'
                   f'    provide_param_guess={guess is not None}\n'
                   f'    test_rolling_window={test_rolling_window}\n' 
-                  f'    plot_frequency     ={plot_step}\n'
-                  f'    plot_kwargs        ={plot_kwargs}\n'
+                  f'    dashboard          ={dashboard is not None}\n'
+                  f'    dashboard_id       ={dashboard_id}\n'
                   f'    print_summary      ={print_summary}\n'
-                  f'    print_progress     ={print_progress}\n')
-            if self.compiled.relaxations:
-                print('Some RDDL operations are non-differentiable, '
-                      'replacing them with differentiable relaxations:')
-                print(self.compiled.summarize_model_relaxations())
-            
+                  f'    print_progress     ={print_progress}\n'
+                  f'    stopping_rule      ={stopping_rule}\n')
+        
+        # ======================================================================
+        # INITIALIZATION OF STATE AND POLICY
+        # ======================================================================
+        
         # compute a batched version of the initial values
         if subs is None:
             subs = self.test_compiled.init_values
@@ -1773,7 +1785,11 @@ class JaxBackpropPlanner:
         else:
             policy_params = guess
             opt_state = self.optimizer.init(policy_params)
-            opt_aux = None
+            opt_aux = {}
+            
+        # ======================================================================
+        # INITIALIZATION OF RUNNING STATISTICS
+        # ======================================================================
         
         # initialize running statistics
         best_params, best_loss, best_grad = policy_params, jnp.inf, jnp.inf
@@ -1783,30 +1799,43 @@ class JaxBackpropPlanner:
         status = JaxPlannerStatus.NORMAL
         is_all_zero_fn = lambda x: np.allclose(x, 0)
         
-        # initialize plot area
-        if plot_step is None or plot_step <= 0 or plt is None:
-            plot = None
-        else:
-            if plot_kwargs is None:
-                plot_kwargs = {}
-            plot = JaxPlannerPlot(self.rddl, self.horizon, **plot_kwargs)
-        xticks, loss_values = [], []
+        # initialize stopping criterion
+        if stopping_rule is not None:
+            stopping_rule.reset()
+            
+        # initialize dash board 
+        if dashboard is not None:
+            dashboard_id = dashboard.register_experiment(
+                dashboard_id, dashboard.get_planner_info(self), 
+                key=dash_key, viz=self.dashboard_viz)
+        
+        # ======================================================================
+        # MAIN TRAINING LOOP BEGINS
+        # ======================================================================
         
-        # training loop
         iters = range(epochs)
         if print_progress:
             iters = tqdm(iters, total=100, position=tqdm_position)
         position_str = '' if tqdm_position is None else f'[{tqdm_position}]'
         
         for it in iters:
+            
+            # ==================================================================
+            # NEXT GRADIENT DESCENT STEP
+            # ==================================================================
+            
             status = JaxPlannerStatus.NORMAL
             
             # update the parameters of the plan
             key, subkey = random.split(key)
-            policy_params, converged, opt_state, opt_aux, \
-            train_loss, train_log = \
+            (policy_params, converged, opt_state, opt_aux, 
+             train_loss, train_log, model_params) = \
                 self.update(subkey, policy_params, policy_hyperparams,
                             train_subs, model_params, opt_state, opt_aux)
+                
+            # ==================================================================
+            # STATUS CHECKS AND LOGGING
+            # ==================================================================
             
             # no progress
             grad_norm_zero, _ = jax.tree_util.tree_flatten(
@@ -1825,12 +1854,11 @@ class JaxBackpropPlanner:
             # numerical error
             if not np.isfinite(train_loss):
                 raise_warning(
-                    f'Aborting JAX planner due to invalid train loss {train_loss}.',
-                    'red')
+                    f'JAX planner aborted due to invalid loss {train_loss}.', 'red')
                 status = JaxPlannerStatus.INVALID_GRADIENT
               
             # evaluate test losses and record best plan so far
-            test_loss, log = self.test_loss(
+            test_loss, (log, model_params_test) = self.test_loss(
                 subkey, policy_params, policy_hyperparams,
                 test_subs, model_params_test)
             test_loss = rolling_test_loss.update(test_loss)
@@ -1839,34 +1867,16 @@ class JaxBackpropPlanner:
                     policy_params, test_loss, train_log['grad']
                 last_iter_improve = it
             
-            # save the plan figure
-            if plot is not None and it % plot_step == 0:
-                xticks.append(it // plot_step)
-                loss_values.append(test_loss.item())
-                action_values = {name: values 
-                                 for (name, values) in log['fluents'].items()
-                                 if name in self.rddl.action_fluents}
-                returns = -np.sum(np.asarray(log['reward']), axis=1)
-                plot.redraw(xticks, loss_values, action_values, returns)
-            
-            # if the progress bar is used
-            elapsed = time.time() - start_time - elapsed_outside_loop
-            if print_progress:
-                iters.n = int(100 * min(1, max(elapsed / train_seconds, it / epochs)))
-                iters.set_description(
-                    f'{position_str} {it:6} it / {-train_loss:14.6f} train / '
-                    f'{-test_loss:14.6f} test / {-best_loss:14.6f} best / '
-                    f'{status.value} status')
-                        
             # reached computation budget
+            elapsed = time.time() - start_time - elapsed_outside_loop
             if elapsed >= train_seconds:
                 status = JaxPlannerStatus.TIME_BUDGET_REACHED
             if it >= epochs - 1:
                 status = JaxPlannerStatus.ITER_BUDGET_REACHED
             
-            # return a callback
-            start_time_outside = time.time()
-            yield {
+            # build a callback
+            progress_percent = int(100 * min(1, max(elapsed / train_seconds, it / epochs)))
+            callback = {
                 'status': status,
                 'iteration': it,
                 'train_return':-train_loss,
@@ -1880,19 +1890,45 @@ class JaxBackpropPlanner:
                 'updates': train_log['updates'],
                 'elapsed_time': elapsed,
                 'key': key,
+                'model_params': model_params,
+                'progress': progress_percent,
+                'train_log': train_log,
                 **log
             }
+            
+            # stopping condition reached
+            if stopping_rule is not None and stopping_rule.monitor(callback):
+                callback['status'] = status = JaxPlannerStatus.STOPPING_RULE_REACHED  
+            
+            # if the progress bar is used
+            if print_progress:
+                iters.n = progress_percent
+                iters.set_description(
+                    f'{position_str} {it:6} it / {-train_loss:14.6f} train / '
+                    f'{-test_loss:14.6f} test / {-best_loss:14.6f} best / '
+                    f'{status.value} status'
+                )
+            
+            # dash-board
+            if dashboard is not None:
+                dashboard.update_experiment(dashboard_id, callback)
+                        
+            # yield the callback
+            start_time_outside = time.time()
+            yield callback
             elapsed_outside_loop += (time.time() - start_time_outside)
             
             # abortion check
-            if status.is_failure():
-                break
-                        
+            if status.is_terminal():
+                break             
+        
+        # ======================================================================
+        # POST-PROCESSING AND CLEANUP
+        # ====================================================================== 
+        
         # release resources
         if print_progress:
             iters.close()
-        if plot is not None:
-            plot.close()
         
         # validate the test return
         if log:
@@ -1918,7 +1954,7 @@ class JaxBackpropPlanner:
                   f'    best_grad_norm={grad_norm}\n'
                   f'    diagnosis: {diagnosis}\n')
     
-    def _perform_diagnosis(self, last_iter_improve, 
+    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_is_zero = np.allclose(max_grad_norm, 0)
@@ -2016,101 +2052,6 @@ class JaxBackpropPlanner:
         return actions      
     
 
-class JaxLineSearchPlanner(JaxBackpropPlanner):
-    '''A class for optimizing an action sequence in the given RDDL MDP using 
-    linear search gradient descent, with the Armijo condition.'''
-    
-    def __init__(self, *args,
-                 decay: float=0.8,
-                 c: float=0.1,
-                 step_max: float=1.0,
-                 step_min: float=1e-6,
-                 **kwargs) -> None:
-        '''Creates a new gradient-based algorithm for optimizing action sequences
-        (plan) in the given RDDL using line search. All arguments are the
-        same as in the parent class, except:
-        
-        :param decay: reduction factor of learning rate per line search iteration
-        :param c: positive coefficient in Armijo condition, should be in (0, 1)
-        :param step_max: initial learning rate for line search
-        :param step_min: minimum possible learning rate (line search halts)
-        '''
-        self.decay = decay
-        self.c = c
-        self.step_max = step_max
-        self.step_min = step_min
-        if 'clip_grad' in kwargs:
-            raise_warning('clip_grad parameter conflicts with '
-                          'line search planner and will be ignored.', 'red')
-            del kwargs['clip_grad']
-        super(JaxLineSearchPlanner, self).__init__(*args, **kwargs)
-        
-    def summarize_hyperparameters(self) -> None:
-        super(JaxLineSearchPlanner, self).summarize_hyperparameters()
-        print(f'linesearch hyper-parameters:\n'
-              f'    decay   ={self.decay}\n'
-              f'    c       ={self.c}\n'
-              f'    lr_range=({self.step_min}, {self.step_max})')
-    
-    def _jax_update(self, loss):
-        optimizer = self.optimizer
-        projection = self.plan.projection
-        decay, c, lrmax, lrmin = self.decay, self.c, self.step_max, self.step_min
-        
-        # initialize the line search routine
-        @jax.jit
-        def _jax_wrapped_line_search_init(key, policy_params, hyperparams,
-                                          subs, model_params):
-            (f, log), grad = jax.value_and_grad(loss, argnums=1, has_aux=True)(
-                key, policy_params, hyperparams, subs, model_params)     
-            gnorm2 = jax.tree_map(lambda x: jnp.sum(jnp.square(x)), grad)
-            gnorm2 = jax.tree_util.tree_reduce(jnp.add, gnorm2)
-            log['grad'] = grad
-            return f, grad, gnorm2, log
-            
-        # compute the next trial solution
-        @jax.jit
-        def _jax_wrapped_line_search_trial(
-                step, grad, key, params, hparams, subs, mparams, state):
-            state.hyperparams['learning_rate'] = step
-            updates, new_state = optimizer.update(grad, state)
-            new_params = optax.apply_updates(params, updates)
-            new_params, _ = projection(new_params, hparams)
-            f_step, _ = loss(key, new_params, hparams, subs, mparams)
-            return f_step, new_params, new_state
-        
-        # main iteration of line search     
-        def _jax_wrapped_plan_update(key, policy_params, hyperparams,
-                                     subs, model_params, opt_state, opt_aux):
-            
-            # initialize the line search
-            f, grad, gnorm2, log = _jax_wrapped_line_search_init(
-                key, policy_params, hyperparams, subs, model_params)            
-            
-            # continue to reduce the learning rate until the Armijo condition holds
-            trials = 0
-            step = lrmax / decay
-            f_step = np.inf
-            best_f, best_step, best_params, best_state = np.inf, None, None, None
-            while (f_step > f - c * step * gnorm2 and step * decay >= lrmin) \
-            or not trials:
-                trials += 1
-                step *= decay
-                f_step, new_params, new_state = _jax_wrapped_line_search_trial(
-                    step, grad, key, policy_params, hyperparams, subs, 
-                    model_params, opt_state)
-                if f_step < best_f:
-                    best_f, best_step, best_params, best_state = \
-                        f_step, step, new_params, new_state
-            
-            log['updates'] = None
-            log['line_search_iters'] = trials
-            log['learning_rate'] = best_step
-            return best_params, True, best_state, best_step, best_f, log
-            
-        return _jax_wrapped_plan_update
-
-
 # ***********************************************************************
 # ALL VERSIONS OF RISK FUNCTIONS
 # 
@@ -2141,7 +2082,7 @@ def cvar_utility(returns: jnp.ndarray, alpha: float) -> float:
     alpha_mask = jax.lax.stop_gradient(
         returns <= jnp.percentile(returns, q=100 * alpha))
     return jnp.sum(returns * alpha_mask) / jnp.sum(alpha_mask)
-    
+
 
 # ***********************************************************************
 # ALL VERSIONS OF CONTROLLERS
@@ -2151,12 +2092,13 @@ def cvar_utility(returns: jnp.ndarray, alpha: float) -> float:
 #
 # ***********************************************************************
 
+
 class JaxOfflineController(BaseAgent):
     '''A container class for a Jax policy trained offline.'''
     
     use_tensor_obs = True
     
-    def __init__(self, planner: JaxBackpropPlanner, 
+    def __init__(self, planner: JaxBackpropPlanner,
                  key: Optional[random.PRNGKey]=None,
                  eval_hyperparams: Optional[Dict[str, Any]]=None,
                  params: Optional[Pytree]=None,
@@ -2186,8 +2128,10 @@ class JaxOfflineController(BaseAgent):
         self.params_given = params is not None
         
         self.step = 0
+        self.callback = None
         if not self.train_on_reset and not self.params_given:
             callback = self.planner.optimize(key=self.key, **self.train_kwargs)
+            self.callback = callback
             params = callback['best_params'] 
         self.params = params  
         
@@ -2202,6 +2146,7 @@ class JaxOfflineController(BaseAgent):
         self.step = 0
         if self.train_on_reset and not self.params_given:
             callback = self.planner.optimize(key=self.key, **self.train_kwargs)
+            self.callback = callback
             self.params = callback['best_params']
 
 
@@ -2211,7 +2156,7 @@ class JaxOnlineController(BaseAgent):
     
     use_tensor_obs = True
     
-    def __init__(self, planner: JaxBackpropPlanner, 
+    def __init__(self, planner: JaxBackpropPlanner,
                  key: Optional[random.PRNGKey]=None,
                  eval_hyperparams: Optional[Dict[str, Any]]=None,
                  warm_start: bool=True,
@@ -2244,7 +2189,9 @@ class JaxOnlineController(BaseAgent):
             key=self.key,
             guess=self.guess,
             subs=state,
-            **self.train_kwargs)
+            **self.train_kwargs
+        )
+        self.callback = callback
         params = callback['best_params']
         self.key, subkey = random.split(self.key)
         actions = planner.get_action(
@@ -2255,4 +2202,5 @@ class JaxOnlineController(BaseAgent):
         
     def reset(self) -> None:
         self.guess = None
+        self.callback = None