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

pyRDDLGym_jax/core/planner.py

@@ -31,22 +31,24 @@ 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]
 Pytree = Any
 
+
 # ***********************************************************************
 # CONFIG FILE MANAGEMENT
 # 
@@ -95,21 +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', {})
-    rounding_name = model_args.get('rounding', 'SoftRounding')
-    rounding_kwargs = model_args.get('rounding_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)
+    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')
@@ -165,6 +171,13 @@ def _load_config(config, args):
     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:
@@ -186,6 +199,7 @@ def load_config_from_string(value: str) -> Tuple[Kwargs, ...]:
     config, args = _parse_config_string(value)
     return _load_config(config, args)
     
+    
 # ***********************************************************************
 # MODEL RELAXATIONS
 # 
@@ -202,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
@@ -237,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:
@@ -298,17 +310,16 @@ 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
 # 
@@ -329,7 +340,7 @@ class JaxPlan:
         self.bounds = None
         
     def summarize_hyperparameters(self) -> None:
-        pass
+        print(self.__str__())
         
     def compile(self, compiled: JaxRDDLCompilerWithGrad,
                 _bounds: Bounds,
@@ -455,21 +466,21 @@ 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,
@@ -820,21 +831,21 @@ 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,
                 horizon: int) -> None:
@@ -1057,6 +1068,7 @@ class JaxDeepReactivePolicy(JaxPlan):
     def guess_next_epoch(self, params: Pytree) -> Pytree:
         return params
     
+    
 # ***********************************************************************
 # ALL VERSIONS OF JAX PLANNER
 # 
@@ -1084,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, 
@@ -1198,14 +1103,15 @@ 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:
@@ -1255,15 +1161,16 @@ class JaxBackpropPlanner:
                  optimizer: Callable[..., optax.GradientTransformation]=optax.rmsprop,
                  optimizer_kwargs: Optional[Kwargs]=None,
                  clip_grad: Optional[float]=None,
-                 noise_grad_eta: float=0.0,
-                 noise_grad_gamma: float=1.0,
-                 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
@@ -1283,9 +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 noise_grad_eta: scale of the gradient noise variance
-        :param noise_grad_gamma: decay rate of the gradient noise variance
-        :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
@@ -1300,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
@@ -1314,31 +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.noise_grad_eta = noise_grad_eta
-        self.noise_grad_gamma = noise_grad_gamma
+        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):
@@ -1371,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 _:
@@ -1394,36 +1308,55 @@ 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'    noise_grad_eta    ={self.noise_grad_eta}\n'
-              f'    noise_grad_gamma  ={self.noise_grad_gamma}\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
@@ -1439,14 +1372,16 @@ r"""
             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,
             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):
@@ -1462,13 +1397,15 @@ r"""
         train_rollouts = self.compiled.compile_rollouts(
             policy=self.plan.train_policy,
             n_steps=self.horizon,
-            n_batch=self.batch_size_train)
+            n_batch=self.batch_size_train
+        )
         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
@@ -1503,14 +1440,16 @@ r"""
         _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
     
@@ -1518,8 +1457,9 @@ r"""
         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, {}
         
@@ -1528,35 +1468,34 @@ r"""
     def _jax_update(self, loss):
         optimizer = self.optimizer
         projection = self.plan.projection
-        
-        # add Gaussian gradient noise per Neelakantan et al., 2016.
-        def _jax_wrapped_gaussian_param_noise(key, grads, sigma):
-            treedef = jax.tree_util.tree_structure(grads)
-            keys_flat = random.split(key, num=treedef.num_leaves)            
-            keys_tree = jax.tree_util.tree_unflatten(treedef, keys_flat)
-            new_grads = jax.tree_map(
-                lambda g, k: g + sigma * random.normal(
-                    key=k, shape=g.shape, dtype=g.dtype),
-                grads,
-                keys_tree
-            )
-            return new_grads
+        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)  
-            sigma = opt_aux.get('noise_sigma', 0.0)
-            grad = _jax_wrapped_gaussian_param_noise(key, grad, sigma)
-            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, opt_aux, loss_val, log
+            return policy_params, converged, opt_state, opt_aux, \
+                loss_val, log, model_params
         
         return jax.jit(_jax_wrapped_plan_update)
             
@@ -1583,7 +1522,6 @@ r"""
         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(
@@ -1637,38 +1575,40 @@ r"""
         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
         
@@ -1683,9 +1623,9 @@ r"""
         
         :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
@@ -1721,8 +1661,8 @@ r"""
     def optimize_generator(self, key: Optional[random.PRNGKey]=None,
                            epochs: int=999999,
                            train_seconds: float=120.,
-                           plot_step: Optional[int]=None,
-                           plot_kwargs: Optional[Kwargs]=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,
@@ -1738,9 +1678,9 @@ r"""
         
         :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
@@ -1760,9 +1700,14 @@ r"""
         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:
@@ -1789,7 +1734,7 @@ r"""
             
         # 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'
@@ -1800,16 +1745,16 @@ r"""
                   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'
                   f'    stopping_rule      ={stopping_rule}\n')
-            if self.compiled.relaxations:
-                print('Some RDDL operations are non-differentiable, '
-                      'they will be approximated as follows:')
-                print(self.compiled.summarize_model_relaxations())
-            
+        
+        # ======================================================================
+        # INITIALIZATION OF STATE AND POLICY
+        # ======================================================================
+        
         # compute a batched version of the initial values
         if subs is None:
             subs = self.test_compiled.init_values
@@ -1841,6 +1786,10 @@ r"""
             policy_params = guess
             opt_state = self.optimizer.init(policy_params)
             opt_aux = {}
+            
+        # ======================================================================
+        # INITIALIZATION OF RUNNING STATISTICS
+        # ======================================================================
         
         # initialize running statistics
         best_params, best_loss, best_grad = policy_params, jnp.inf, jnp.inf
@@ -1854,35 +1803,39 @@ r"""
         if stopping_rule is not None:
             stopping_rule.reset()
             
-        # 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 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:
-            status = JaxPlannerStatus.NORMAL
             
-            # gradient noise schedule
-            noise_var = self.noise_grad_eta / (1. + it) ** self.noise_grad_gamma
-            noise_sigma = np.sqrt(noise_var)
-            opt_aux['noise_sigma'] = noise_sigma
+            # ==================================================================
+            # 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(
@@ -1901,12 +1854,11 @@ r"""
             # 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)
@@ -1915,32 +1867,15 @@ r"""
                     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
+            # build a callback
+            progress_percent = int(100 * min(1, max(elapsed / train_seconds, it / epochs)))
             callback = {
                 'status': status,
                 'iteration': it,
@@ -1952,29 +1887,48 @@ r"""
                 'last_iteration_improved': last_iter_improve,
                 'grad': train_log['grad'],
                 'best_grad': best_grad,
-                'noise_sigma': noise_sigma,
                 '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
-            
-            # stopping condition reached
-            if stopping_rule is not None and stopping_rule.monitor(callback):
-                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:
@@ -2098,101 +2052,6 @@ r"""
         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
-            opt_aux['best_step'] = best_step
-            return best_params, True, best_state, opt_aux, best_f, log
-            
-        return _jax_wrapped_plan_update
-
 # ***********************************************************************
 # ALL VERSIONS OF RISK FUNCTIONS
 # 
@@ -2224,6 +2083,7 @@ def cvar_utility(returns: jnp.ndarray, alpha: float) -> float:
         returns <= jnp.percentile(returns, q=100 * alpha))
     return jnp.sum(returns * alpha_mask) / jnp.sum(alpha_mask)
 
+
 # ***********************************************************************
 # ALL VERSIONS OF CONTROLLERS
 # 
@@ -2268,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  
         
@@ -2284,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']
 
 
@@ -2326,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(
@@ -2337,4 +2202,5 @@ class JaxOnlineController(BaseAgent):
         
     def reset(self) -> None:
         self.guess = None
+        self.callback = None