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

pyRDDLGym_jax/core/planner.py

@@ -2,54 +2,50 @@ from ast import literal_eval
 from collections import deque
 import configparser
 from enum import Enum
+import os
+import sys
+import time
+import traceback
+from typing import Any, Callable, Dict, Generator, Optional, Set, Sequence, Tuple, Union
+
 import haiku as hk
 import jax
+import jax.nn.initializers as initializers
 import jax.numpy as jnp
 import jax.random as random
-import jax.nn.initializers as initializers
 import numpy as np
 import optax
-import os
-import sys
 import termcolor
-import time
-import traceback
 from tqdm import tqdm
-from typing import Any, Callable, Dict, Generator, Optional, Set, Sequence, Tuple, Union
 
-Activation = Callable[[jnp.ndarray], jnp.ndarray]
-Bounds = Dict[str, Tuple[np.ndarray, np.ndarray]]
-Kwargs = Dict[str, Any]
-Pytree = Any
-
-from pyRDDLGym.core.debug.exception import raise_warning
-
-from pyRDDLGym_jax import __version__
-
-# try to import matplotlib, if failed then skip plotting
-try:
-    import matplotlib
-    import matplotlib.pyplot as plt
-    matplotlib.use('TkAgg')
-except Exception:
-    raise_warning('failed to import matplotlib: '
-                  'plotting functionality will be disabled.', 'red')
-    traceback.print_exc()
-    plt = None
-            
 from pyRDDLGym.core.compiler.model import RDDLPlanningModel, RDDLLiftedModel
 from pyRDDLGym.core.debug.logger import Logger
 from pyRDDLGym.core.debug.exception import (
+    raise_warning,
     RDDLNotImplementedError,
     RDDLUndefinedVariableError,
     RDDLTypeError
 )
 from pyRDDLGym.core.policy import BaseAgent
 
-from pyRDDLGym_jax.core.compiler import JaxRDDLCompiler
+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
 
+# try to import matplotlib, if failed then skip plotting
+try:
+    import matplotlib.pyplot as plt
+except Exception:
+    raise_warning('failed to import matplotlib: '
+                  'plotting functionality will be disabled.', 'red')
+    traceback.print_exc()
+    plt = None
+    
+Activation = Callable[[jnp.ndarray], jnp.ndarray]
+Bounds = Dict[str, Tuple[np.ndarray, np.ndarray]]
+Kwargs = Dict[str, Any]
+Pytree = Any
 
 # ***********************************************************************
 # CONFIG FILE MANAGEMENT
@@ -60,6 +56,7 @@ from pyRDDLGym_jax.core.logic import FuzzyLogic
 #
 # ***********************************************************************
 
+
 def _parse_config_file(path: str):
     if not os.path.isfile(path):
         raise FileNotFoundError(f'File {path} does not exist.')
@@ -104,9 +101,15 @@ def _load_config(config, args):
     comp_kwargs = model_args.get('complement_kwargs', {})
     compare_name = model_args.get('comparison', 'SigmoidComparison')
     compare_kwargs = model_args.get('comparison_kwargs', {})
+    sampling_name = model_args.get('sampling', 'GumbelSoftmax')
+    sampling_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)
     
     # read the policy settings
     plan_method = planner_args.pop('method')
@@ -157,11 +160,18 @@ 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)
     
+    # 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 +185,6 @@ 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
@@ -184,18 +193,6 @@ def load_config_from_string(value: str) -> Tuple[Kwargs, ...]:
 #
 # ***********************************************************************
 
-def _function_discrete_approx_named(logic):
-    jax_discrete, jax_param = logic.discrete()
-
-    def _jax_wrapped_discrete_calc_approx(key, prob, params):
-        sample = jax_discrete(key, prob, params)
-        out_of_bounds = jnp.logical_not(jnp.logical_and(
-            jnp.all(prob >= 0),
-            jnp.allclose(jnp.sum(prob, axis=-1), 1.0)))
-        return sample, out_of_bounds
-    
-    return _jax_wrapped_discrete_calc_approx, jax_param
-
 
 class JaxRDDLCompilerWithGrad(JaxRDDLCompiler):
     '''Compiles a RDDL AST representation to an equivalent JAX representation. 
@@ -271,7 +268,9 @@ class JaxRDDLCompilerWithGrad(JaxRDDLCompiler):
         self.IF_HELPER = logic.control_if()
         self.SWITCH_HELPER = logic.control_switch()
         self.BERNOULLI_HELPER = logic.bernoulli()
-        self.DISCRETE_HELPER = _function_discrete_approx_named(logic)
+        self.DISCRETE_HELPER = logic.discrete()
+        self.POISSON_HELPER = logic.poisson()
+        self.GEOMETRIC_HELPER = logic.geometric()
         
     def _jax_stop_grad(self, jax_expr):
         
@@ -309,7 +308,6 @@ class JaxRDDLCompilerWithGrad(JaxRDDLCompiler):
         arg, = expr.args
         arg = self._jax(arg, info)
         return arg
-    
 
 # ***********************************************************************
 # ALL VERSIONS OF JAX PLANS
@@ -319,6 +317,7 @@ class JaxRDDLCompilerWithGrad(JaxRDDLCompiler):
 #
 # ***********************************************************************
 
+
 class JaxPlan:
     '''Base class for all JAX policy representations.'''
     
@@ -373,7 +372,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():
@@ -469,10 +468,11 @@ class JaxStraightLinePlan(JaxPlan):
               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}')
+              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
         
@@ -513,7 +513,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
             
@@ -522,14 +522,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
@@ -789,7 +788,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:
@@ -837,7 +836,7 @@ class JaxDeepReactivePolicy(JaxPlan):
               f'    wrap_non_bool       ={self._wrap_non_bool}')
     
     def compile(self, compiled: JaxRDDLCompilerWithGrad,
-                _bounds: Bounds, 
+                _bounds: Bounds,
                 horizon: int) -> None:
         rddl = compiled.rddl
         
@@ -890,7 +889,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
@@ -915,8 +914,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)
@@ -926,7 +925,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)
@@ -959,14 +958,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
@@ -1058,7 +1056,6 @@ class JaxDeepReactivePolicy(JaxPlan):
         
     def guess_next_epoch(self, params: Pytree) -> Pytree:
         return params
-
     
 # ***********************************************************************
 # ALL VERSIONS OF JAX PLANNER
@@ -1068,6 +1065,7 @@ class JaxDeepReactivePolicy(JaxPlan):
 #
 # ***********************************************************************
 
+
 class RollingMean:
     '''Maintains an estimate of the rolling mean of a stream of real-valued 
     observations.'''
@@ -1089,7 +1087,7 @@ class RollingMean:
 class JaxPlannerPlot:
     '''Supports plotting and visualization of a JAX policy in real time.'''
     
-    def __init__(self, rddl: RDDLPlanningModel, horizon: int, 
+    def __init__(self, rddl: RDDLPlanningModel, horizon: int,
                  show_violin: bool=True, show_action: bool=True) -> None:
         '''Creates a new planner visualizer.
         
@@ -1137,7 +1135,7 @@ class JaxPlannerPlot:
                 for dim in rddl.object_counts(rddl.variable_params[name]):
                     action_dim *= dim     
                 action_plot = ax.pcolormesh(
-                    np.zeros((action_dim, horizon)), 
+                    np.zeros((action_dim, horizon)),
                     cmap='seismic', vmin=vmin, vmax=vmax)
                 ax.set_aspect('auto')        
                 ax.set_xlabel('decision epoch')
@@ -1210,6 +1208,39 @@ class JaxPlannerStatus(Enum):
         return self.value >= 3
 
 
+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.'''
@@ -1224,6 +1255,8 @@ 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(),
                  use_symlog_reward: bool=False,
                  utility: Union[Callable[[jnp.ndarray], float], str]='mean',
@@ -1250,6 +1283,8 @@ 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
         operations to their differentiable counterparts 
         :param use_symlog_reward: whether to use the symlog transform on the 
@@ -1284,6 +1319,8 @@ class JaxBackpropPlanner:
             optimizer_kwargs = {'learning_rate': 0.1}
         self._optimizer_kwargs = optimizer_kwargs
         self.clip_grad = clip_grad
+        self.noise_grad_eta = noise_grad_eta
+        self.noise_grad_gamma = noise_grad_gamma
         
         # set optimizer
         try:
@@ -1348,8 +1385,18 @@ class JaxBackpropPlanner:
                 map(str, jax._src.xla_bridge.devices())).replace('\n', '')
         except Exception as _:
             devices_short = 'N/A'
+        LOGO = \
+r"""
+   __   ______   __  __   ______  __       ______   __   __    
+  /\ \ /\  __ \ /\_\_\_\ /\  == \/\ \     /\  __ \ /\ "-.\ \   
+ _\_\ \\ \  __ \\/_/\_\/_\ \  _-/\ \ \____\ \  __ \\ \ \-.  \  
+/\_____\\ \_\ \_\ /\_\/\_\\ \_\   \ \_____\\ \_\ \_\\ \_\\"\_\ 
+\/_____/ \/_/\/_/ \/_/\/_/ \/_/    \/_____/ \/_/\/_/ \/_/ \/_/ 
+"""
+                   
         print('\n'
-              f'JAX Planner version {__version__}\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__}, '
@@ -1371,6 +1418,8 @@ class JaxBackpropPlanner:
               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()
@@ -1395,7 +1444,7 @@ class JaxBackpropPlanner:
         
         # Jax compilation of the exact RDDL for testing
         self.test_compiled = JaxRDDLCompiler(
-            rddl=rddl, 
+            rddl=rddl,
             logger=self.logger,
             use64bit=self.use64bit)
         self.test_compiled.compile(log_jax_expr=True, heading='EXACT MODEL')
@@ -1472,7 +1521,7 @@ class JaxBackpropPlanner:
         def _jax_wrapped_init_policy(key, hyperparams, subs):
             policy_params = init(key, hyperparams, subs)
             opt_state = optimizer.init(policy_params)
-            return policy_params, opt_state, None
+            return policy_params, opt_state, {}
         
         return _jax_wrapped_init_policy
         
@@ -1480,6 +1529,19 @@ class JaxBackpropPlanner:
         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
+        
         # 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,
@@ -1487,12 +1549,14 @@ class JaxBackpropPlanner:
             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) 
             policy_params = optax.apply_updates(policy_params, updates)
             policy_params, converged = projection(policy_params, 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
         
         return jax.jit(_jax_wrapped_plan_update)
             
@@ -1523,7 +1587,7 @@ class JaxBackpropPlanner:
         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]:
@@ -1575,7 +1639,7 @@ class JaxBackpropPlanner:
         @jax.jit
         def _loss_with_key(key, params_1d):
             policy_params = unravel_fn(params_1d)
-            loss_val, _ = loss_fn(key, policy_params, policy_hyperparams, 
+            loss_val, _ = loss_fn(key, policy_params, policy_hyperparams,
                                   train_subs, model_params)
             return loss_val
         
@@ -1583,7 +1647,7 @@ class JaxBackpropPlanner:
         def _grad_with_key(key, params_1d):
             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, 
+            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
@@ -1632,6 +1696,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,13 +1722,14 @@ class JaxBackpropPlanner:
                            epochs: int=999999,
                            train_seconds: float=120.,
                            plot_step: Optional[int]=None,
-                           plot_kwargs: Optional[Dict[str, Any]]=None,
+                           plot_kwargs: Optional[Kwargs]=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. 
@@ -1685,6 +1751,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
@@ -1711,6 +1778,14 @@ class JaxBackpropPlanner:
             hyperparam_value = float(policy_hyperparams)
             policy_hyperparams = {action: hyperparam_value
                                   for action in self.rddl.action_fluents}
+        
+        # fill in missing entries
+        elif isinstance(policy_hyperparams, dict):
+            for action in self.rddl.action_fluents:
+                if action not in policy_hyperparams:
+                    raise_warning(f'policy_hyperparams[{action}] is not set, '
+                                  'setting 1.0 which could be suboptimal.')
+                    policy_hyperparams[action] = 1.0
             
         # print summary of parameters:
         if print_summary:
@@ -1728,10 +1803,11 @@ class JaxBackpropPlanner:
                   f'    plot_frequency     ={plot_step}\n'
                   f'    plot_kwargs        ={plot_kwargs}\n'
                   f'    print_summary      ={print_summary}\n'
-                  f'    print_progress     ={print_progress}\n')
+                  f'    print_progress     ={print_progress}\n'
+                  f'    stopping_rule      ={stopping_rule}\n')
             if self.compiled.relaxations:
                 print('Some RDDL operations are non-differentiable, '
-                      'replacing them with differentiable relaxations:')
+                      'they will be approximated as follows:')
                 print(self.compiled.summarize_model_relaxations())
             
         # compute a batched version of the initial values
@@ -1764,7 +1840,7 @@ class JaxBackpropPlanner:
         else:
             policy_params = guess
             opt_state = self.optimizer.init(policy_params)
-            opt_aux = None
+            opt_aux = {}
         
         # initialize running statistics
         best_params, best_loss, best_grad = policy_params, jnp.inf, jnp.inf
@@ -1772,7 +1848,12 @@ class JaxBackpropPlanner:
         rolling_test_loss = RollingMean(test_rolling_window)
         log = {}
         status = JaxPlannerStatus.NORMAL
+        is_all_zero_fn = lambda x: np.allclose(x, 0)
         
+        # initialize stopping criterion
+        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
@@ -1786,10 +1867,16 @@ class JaxBackpropPlanner:
         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
+            
             # update the parameters of the plan
             key, subkey = random.split(key)
             policy_params, converged, opt_state, opt_aux, \
@@ -1799,7 +1886,7 @@ class JaxBackpropPlanner:
             
             # no progress
             grad_norm_zero, _ = jax.tree_util.tree_flatten(
-                jax.tree_map(lambda x: np.allclose(x, 0), train_log['grad']))
+                jax.tree_map(is_all_zero_fn, train_log['grad']))
             if np.all(grad_norm_zero):
                 status = JaxPlannerStatus.NO_PROGRESS
              
@@ -1843,8 +1930,9 @@ class JaxBackpropPlanner:
             if print_progress:
                 iters.n = int(100 * min(1, max(elapsed / train_seconds, it / epochs)))
                 iters.set_description(
-                    f'[{tqdm_position}] {it:6} it / {-train_loss:14.6f} train / '
-                    f'{-test_loss:14.6f} test / {-best_loss:14.6f} best')
+                    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
             if elapsed >= train_seconds:
@@ -1853,8 +1941,7 @@ class JaxBackpropPlanner:
                 status = JaxPlannerStatus.ITER_BUDGET_REACHED
             
             # return a callback
-            start_time_outside = time.time()
-            yield {
+            callback = {
                 'status': status,
                 'iteration': it,
                 'train_return':-train_loss,
@@ -1865,16 +1952,23 @@ class JaxBackpropPlanner:
                 '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,
                 **log
             }
+            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                
                         
         # release resources
         if print_progress:
@@ -1904,9 +1998,9 @@ class JaxBackpropPlanner:
                   f'    iterations    ={it}\n'
                   f'    best_objective={-best_loss}\n'
                   f'    best_grad_norm={grad_norm}\n'
-                  f'diagnosis: {diagnosis}\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)
@@ -2085,7 +2179,7 @@ class JaxLineSearchPlanner(JaxBackpropPlanner):
                 trials += 1
                 step *= decay
                 f_step, new_params, new_state = _jax_wrapped_line_search_trial(
-                    step, grad, key, policy_params, hyperparams, subs, 
+                    step, grad, key, policy_params, hyperparams, subs,
                     model_params, opt_state)
                 if f_step < best_f:
                     best_f, best_step, best_params, best_state = \
@@ -2094,11 +2188,11 @@ class JaxLineSearchPlanner(JaxBackpropPlanner):
             log['updates'] = None
             log['line_search_iters'] = trials
             log['learning_rate'] = best_step
-            return best_params, True, best_state, best_step, best_f, log
+            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
 # 
@@ -2116,7 +2210,7 @@ class JaxLineSearchPlanner(JaxBackpropPlanner):
 @jax.jit
 def entropic_utility(returns: jnp.ndarray, beta: float) -> float:
     return (-1.0 / beta) * jax.scipy.special.logsumexp(
-            -beta * returns, b=1.0 / returns.size)
+        -beta * returns, b=1.0 / returns.size)
 
 
 @jax.jit
@@ -2129,7 +2223,6 @@ 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
@@ -2139,12 +2232,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,
@@ -2199,7 +2293,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,