pyRDDLGym-jax 0.1__py3-none-any.whl → 0.2__py3-none-any.whl

pyRDDLGym_jax/core/planner.py

@@ -1,6 +1,9 @@
+__version__ = '0.2'
+
 from ast import literal_eval
 from collections import deque
 import configparser
+from enum import Enum
 import haiku as hk
 import jax
 import jax.numpy as jnp
@@ -13,11 +16,28 @@ import sys
 import termcolor
 import time
 from tqdm import tqdm
-from typing import Callable, Dict, Generator, Set, Sequence, Tuple
+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
 
+# try to import matplotlib, if failed then skip plotting
+try:
+    import matplotlib
+    import matplotlib.pyplot as plt
+    matplotlib.use('TkAgg')
+except Exception:
+    raise_warning('matplotlib is not installed, '
+                  'plotting functionality is disabled.', 'red')
+    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
@@ -37,6 +57,7 @@ from pyRDDLGym_jax.core.logic import FuzzyLogic
 # - instantiate planner
 #
 # ***********************************************************************
+
 def _parse_config_file(path: str):
     if not os.path.isfile(path):
         raise FileNotFoundError(f'File {path} does not exist.')
@@ -59,51 +80,94 @@ def _parse_config_string(value: str):
     return config, args
 
 
+def _getattr_any(packages, item):
+    for package in packages:
+        loaded = getattr(package, item, None)
+        if loaded is not None:
+            return loaded
+    return None
+
+
 def _load_config(config, args):
     model_args = {k: args[k] for (k, _) in config.items('Model')}
     planner_args = {k: args[k] for (k, _) in config.items('Optimizer')}
     train_args = {k: args[k] for (k, _) in config.items('Training')}    
     
-    train_args['key'] = jax.random.PRNGKey(train_args['key'])
-    
     # read the model settings
-    tnorm_name = model_args['tnorm']
-    tnorm_kwargs = model_args['tnorm_kwargs']
-    logic_name = model_args['logic']
-    logic_kwargs = model_args['logic_kwargs']
+    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', {})
     logic_kwargs['tnorm'] = getattr(logic, tnorm_name)(**tnorm_kwargs)
-    planner_args['logic'] = getattr(logic, logic_name)(**logic_kwargs)
+    logic_kwargs['complement'] = getattr(logic, comp_name)(**comp_kwargs)
+    logic_kwargs['comparison'] = getattr(logic, compare_name)(**compare_kwargs)
     
-    # read the optimizer settings
+    # read the policy settings
     plan_method = planner_args.pop('method')
     plan_kwargs = planner_args.pop('method_kwargs', {})  
     
-    if 'initializer' in plan_kwargs:  # weight initialization
-        init_name = plan_kwargs['initializer']
-        init_class = getattr(initializers, init_name)
-        init_kwargs = plan_kwargs.pop('initializer_kwargs', {})
-        try: 
-            plan_kwargs['initializer'] = init_class(**init_kwargs)
-        except:
-            raise_warning(f'ignoring arguments for initializer <{init_name}>')
-            plan_kwargs['initializer'] = init_class
-               
-    if 'activation' in plan_kwargs:  # activation function
-        plan_kwargs['activation'] = getattr(jax.nn, plan_kwargs['activation'])
+    # policy initialization
+    plan_initializer = plan_kwargs.get('initializer', None)
+    if plan_initializer is not None:
+        initializer = _getattr_any(packages=[initializers], item=plan_initializer)
+        if initializer is None:
+            raise_warning(
+                f'Ignoring invalid initializer <{plan_initializer}>.', 'red')
+            del plan_kwargs['initializer']
+        else:
+            init_kwargs = plan_kwargs.pop('initializer_kwargs', {})
+            try: 
+                plan_kwargs['initializer'] = initializer(**init_kwargs)
+            except Exception as _:
+                raise_warning(
+                    f'Ignoring invalid initializer_kwargs <{init_kwargs}>.', 'red')
+                plan_kwargs['initializer'] = initializer
     
+    # policy activation
+    plan_activation = plan_kwargs.get('activation', None)
+    if plan_activation is not None:
+        activation = _getattr_any(packages=[jax.nn, jax.numpy], item=plan_activation)
+        if activation is None:
+            raise_warning(
+                f'Ignoring invalid activation <{plan_activation}>.', 'red')
+            del plan_kwargs['activation']
+        else:
+            plan_kwargs['activation'] = activation
+    
+    # read the planner settings
+    planner_args['logic'] = getattr(logic, logic_name)(**logic_kwargs)
     planner_args['plan'] = getattr(sys.modules[__name__], plan_method)(**plan_kwargs)
-    planner_args['optimizer'] = getattr(optax, planner_args['optimizer'])
+    
+    # planner optimizer
+    planner_optimizer = planner_args.get('optimizer', None)
+    if planner_optimizer is not None:
+        optimizer = _getattr_any(packages=[optax], item=planner_optimizer)
+        if optimizer is None:
+            raise_warning(
+                f'Ignoring invalid optimizer <{planner_optimizer}>.', 'red')
+            del planner_args['optimizer']
+        else:
+            planner_args['optimizer'] = optimizer
+        
+    # read the optimize call settings
+    planner_key = train_args.get('key', None)
+    if planner_key is not None:
+        train_args['key'] = random.PRNGKey(planner_key)
     
     return planner_args, plan_kwargs, train_args
 
 
-def load_config(path: str) -> Tuple[Dict[str, object], ...]:
+def load_config(path: str) -> Tuple[Kwargs, ...]:
     '''Loads a config file at the specified file path.'''
     config, args = _parse_config_file(path)
     return _load_config(config, args)
 
 
-def load_config_from_string(value: str) -> Tuple[Dict[str, object], ...]:
+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)
@@ -115,6 +179,20 @@ def load_config_from_string(value: str) -> Tuple[Dict[str, object], ...]:
 # - replace discrete ops in state dynamics/reward with differentiable ones
 #
 # ***********************************************************************
+
+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. 
     Unlike its parent class, this class treats all fluents as real-valued, and
@@ -124,7 +202,7 @@ class JaxRDDLCompilerWithGrad(JaxRDDLCompiler):
     
     def __init__(self, *args,
                  logic: FuzzyLogic=FuzzyLogic(),
-                 cpfs_without_grad: Set=set(),
+                 cpfs_without_grad: Optional[Set[str]]=None,
                  **kwargs) -> None:
         '''Creates a new RDDL to Jax compiler, where operations that are not
         differentiable are converted to approximate forms that have defined 
@@ -140,27 +218,30 @@ class JaxRDDLCompilerWithGrad(JaxRDDLCompiler):
         '''
         super(JaxRDDLCompilerWithGrad, self).__init__(*args, **kwargs)
         self.logic = logic
+        self.logic.set_use64bit(self.use64bit)
+        if cpfs_without_grad is None:
+            cpfs_without_grad = set()
         self.cpfs_without_grad = cpfs_without_grad
         
         # actions and CPFs must be continuous
-        raise_warning(f'Initial values of pvariables will be cast to real.')   
+        raise_warning('Initial values of pvariables will be cast to real.')   
         for (var, values) in self.init_values.items():
             self.init_values[var] = np.asarray(values, dtype=self.REAL) 
         
         # overwrite basic operations with fuzzy ones
         self.RELATIONAL_OPS = {
-            '>=': logic.greaterEqual(),
-            '<=': logic.lessEqual(),
+            '>=': logic.greater_equal(),
+            '<=': logic.less_equal(),
             '<': logic.less(),
             '>': logic.greater(),
             '==': logic.equal(),
-            '~=': logic.notEqual()
+            '~=': logic.not_equal()
         }
-        self.LOGICAL_NOT = logic.Not()
+        self.LOGICAL_NOT = logic.logical_not()
         self.LOGICAL_OPS = {
-            '^': logic.And(),
-            '&': logic.And(),
-            '|': logic.Or(),
+            '^': logic.logical_and(),
+            '&': logic.logical_and(),
+            '|': logic.logical_or(),
             '~': logic.xor(),
             '=>': logic.implies(),
             '<=>': logic.equiv()
@@ -169,15 +250,19 @@ class JaxRDDLCompilerWithGrad(JaxRDDLCompiler):
         self.AGGREGATION_OPS['exists'] = logic.exists()
         self.AGGREGATION_OPS['argmin'] = logic.argmin()
         self.AGGREGATION_OPS['argmax'] = logic.argmax()
-        self.KNOWN_UNARY['sgn'] = logic.signum()
+        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.floorDiv()
+        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 = _function_discrete_approx_named(logic)
+        
     def _jax_stop_grad(self, jax_expr):
         
         def _jax_wrapped_stop_grad(x, params, key):
@@ -199,35 +284,13 @@ class JaxRDDLCompilerWithGrad(JaxRDDLCompiler):
                     jax_cpfs[cpf] = self._jax_stop_grad(jax_cpfs[cpf])
         return jax_cpfs
     
-    def _jax_if_helper(self):
-        return self.logic.If()
-    
-    def _jax_switch_helper(self):
-        return self.logic.Switch()
-        
     def _jax_kron(self, expr, info):
         if self.logic.verbose:
             raise_warning('KronDelta will be ignored.')            
-                       
         arg, = expr.args
         arg = self._jax(arg, info)
         return arg
     
-    def _jax_bernoulli_helper(self):
-        return self.logic.bernoulli()
-    
-    def _jax_discrete_helper(self):
-        jax_discrete, jax_param = self.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
-
 
 # ***********************************************************************
 # ALL VERSIONS OF JAX PLANS
@@ -236,6 +299,7 @@ class JaxRDDLCompilerWithGrad(JaxRDDLCompiler):
 # - deep reactive policy
 #
 # ***********************************************************************
+
 class JaxPlan:
     '''Base class for all JAX policy representations.'''
     
@@ -245,15 +309,15 @@ class JaxPlan:
         self._test_policy = None
         self._projection = None
     
-    def summarize_hyperparameters(self):
+    def summarize_hyperparameters(self) -> None:
         pass
         
     def compile(self, compiled: JaxRDDLCompilerWithGrad,
-                _bounds: Dict,
+                _bounds: Bounds,
                 horizon: int) -> None:
         raise NotImplementedError
     
-    def guess_next_epoch(self, params: Dict) -> Dict:
+    def guess_next_epoch(self, params: Pytree) -> Pytree:
         raise NotImplementedError
     
     @property
@@ -289,7 +353,8 @@ class JaxPlan:
         self._projection = value
     
     def _calculate_action_info(self, compiled: JaxRDDLCompilerWithGrad,
-                               user_bounds: Dict[str, object], horizon: int):
+                               user_bounds: Bounds, 
+                               horizon: int):
         shapes, bounds, bounds_safe, cond_lists = {}, {}, {}, {}
         for (name, prange) in compiled.rddl.variable_ranges.items():
             if compiled.rddl.variable_types[name] != 'action-fluent':
@@ -309,8 +374,8 @@ class JaxPlan:
             else:
                 lower, upper = compiled.constraints.bounds[name]
                 lower, upper = user_bounds.get(name, (lower, upper))
-                lower = np.asarray(lower, dtype=np.float32)
-                upper = np.asarray(upper, dtype=np.float32)
+                lower = np.asarray(lower, dtype=compiled.REAL)
+                upper = np.asarray(upper, dtype=compiled.REAL)
                 lower_finite = np.isfinite(lower)
                 upper_finite = np.isfinite(upper)
                 bounds_safe[name] = (np.where(lower_finite, lower, 0.0),
@@ -336,7 +401,7 @@ class JaxStraightLinePlan(JaxPlan):
     
     def __init__(self, initializer: initializers.Initializer=initializers.normal(),
                  wrap_sigmoid: bool=True,
-                 min_action_prob: float=1e-5,
+                 min_action_prob: float=1e-6,
                  wrap_non_bool: bool=False,
                  wrap_softmax: bool=False,
                  use_new_projection: bool=False,
@@ -371,7 +436,7 @@ class JaxStraightLinePlan(JaxPlan):
         self._use_new_projection = use_new_projection
         self._max_constraint_iter = max_constraint_iter
         
-    def summarize_hyperparameters(self):
+    def summarize_hyperparameters(self) -> None:
         print(f'policy hyper-parameters:\n'
               f'    initializer          ={type(self._initializer_base).__name__}\n'
               f'constraint-sat strategy (simple):\n'
@@ -383,7 +448,8 @@ class JaxStraightLinePlan(JaxPlan):
               f'    use_new_projection   ={self._use_new_projection}')
     
     def compile(self, compiled: JaxRDDLCompilerWithGrad,
-                _bounds: Dict, horizon: int) -> None:
+                _bounds: Bounds, 
+                horizon: int) -> None:
         rddl = compiled.rddl
         
         # calculate the correct action box bounds
@@ -423,7 +489,7 @@ class JaxStraightLinePlan(JaxPlan):
         def _jax_bool_action_to_param(var, action, hyperparams):
             if wrap_sigmoid:
                 weight = hyperparams[var]
-                return (-1.0 / weight) * jnp.log1p(1.0 / action - 2.0)
+                return (-1.0 / weight) * jnp.log(1.0 / action - 1.0)
             else:
                 return action
             
@@ -506,7 +572,7 @@ class JaxStraightLinePlan(JaxPlan):
         def _jax_wrapped_slp_predict_test(key, params, hyperparams, step, subs):
             actions = {}
             for (var, param) in params.items():
-                action = jnp.asarray(param[step, ...])
+                action = jnp.asarray(param[step, ...], dtype=compiled.REAL)
                 if var == bool_key:
                     output = jax.nn.softmax(action)
                     bool_actions = _jax_unstack_bool_from_softmax(output)
@@ -688,7 +754,7 @@ class JaxStraightLinePlan(JaxPlan):
         # "progress" the plan one step forward and set last action to second-last
         return jnp.append(param[1:, ...], param[-1:, ...], axis=0)
 
-    def guess_next_epoch(self, params: Dict) -> Dict:
+    def guess_next_epoch(self, params: Pytree) -> Pytree:
         next_fn = JaxStraightLinePlan._guess_next_epoch
         return jax.tree_map(next_fn, params)
 
@@ -696,10 +762,12 @@ class JaxStraightLinePlan(JaxPlan):
 class JaxDeepReactivePolicy(JaxPlan):
     '''A deep reactive policy network implementation in JAX.'''
     
-    def __init__(self, topology: Sequence[int],
-                 activation: Callable=jax.nn.relu,
+    def __init__(self, topology: Optional[Sequence[int]]=None,
+                 activation: Activation=jnp.tanh,
                  initializer: hk.initializers.Initializer=hk.initializers.VarianceScaling(scale=2.0),
-                 normalize: bool=True) -> None:
+                 normalize: bool=True, 
+                 normalizer_kwargs: Optional[Kwargs]=None,
+                 wrap_non_bool: bool=False) -> None:
         '''Creates a new deep reactive policy in JAX.
         
         :param neurons: sequence consisting of the number of neurons in each
@@ -707,23 +775,39 @@ class JaxDeepReactivePolicy(JaxPlan):
         :param activation: function to apply after each layer of the policy
         :param initializer: weight initialization
         :param normalize: whether to apply layer norm to the inputs
+        :param normalizer_kwargs: if normalize is True, apply additional arguments
+        to layer norm
+        :param wrap_non_bool: whether to wrap real or int action fluent parameters
+        with non-linearity (e.g. sigmoid or ELU) to satisfy box constraints
         '''
         super(JaxDeepReactivePolicy, self).__init__()
+        if topology is None:
+            topology = [128, 64]
         self._topology = topology
         self._activations = [activation for _ in topology]
         self._initializer_base = initializer
         self._initializer = initializer
         self._normalize = normalize
+        if normalizer_kwargs is None:
+            normalizer_kwargs = {
+                'create_offset': True, 'create_scale': True, 
+                'name': 'input_norm'
+            }
+        self._normalizer_kwargs = normalizer_kwargs
+        self._wrap_non_bool = wrap_non_bool
             
-    def summarize_hyperparameters(self):
+    def summarize_hyperparameters(self) -> None:
         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_layer_norm={self._normalize}')
+              f'    apply_layer_norm={self._normalize}\n'
+              f'    layer_norm_args ={self._normalizer_kwargs}\n'
+              f'    wrap_non_bool   ={self._wrap_non_bool}')
     
     def compile(self, compiled: JaxRDDLCompilerWithGrad,
-                _bounds: Dict, horizon: int) -> None:
+                _bounds: Bounds, 
+                horizon: int) -> None:
         rddl = compiled.rddl
         
         # calculate the correct action box bounds
@@ -751,6 +835,7 @@ class JaxDeepReactivePolicy(JaxPlan):
                    
         ranges = rddl.variable_ranges
         normalize = self._normalize
+        wrap_non_bool = self._wrap_non_bool
         init = self._initializer
         layers = list(enumerate(zip(self._topology, self._activations)))
         layer_sizes = {var: np.prod(shape, dtype=int) 
@@ -763,9 +848,7 @@ class JaxDeepReactivePolicy(JaxPlan):
             # apply layer norm
             if normalize:
                 normalizer = hk.LayerNorm(
-                    axis=-1, param_axis=-1,
-                    create_offset=True, create_scale=True,
-                    name='input_norm')
+                    axis=-1, param_axis=-1, **self._normalizer_kwargs)
                 state = normalizer(state)
             
             # feed state vector through hidden layers
@@ -789,16 +872,19 @@ class JaxDeepReactivePolicy(JaxPlan):
                     if not use_constraint_satisfaction:
                         actions[var] = jax.nn.sigmoid(output)
                 else:
-                    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
-                        ]
-                    )
+                    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
+                            ]
+                        )
+                    else:
+                        action = output
                     actions[var] = action
             
             # for constraint satisfaction wrap bool actions with softmax
@@ -826,12 +912,17 @@ class JaxDeepReactivePolicy(JaxPlan):
                     actions[name] = action
                     start += size
             return actions
-                
+        
+        if rddl.observ_fluents:
+            observed_vars = rddl.observ_fluents
+        else:
+            observed_vars = rddl.state_fluents
+             
         # state is concatenated into single tensor
         def _jax_wrapped_subs_to_state(subs):
             subs = {var: value
                     for (var, value) in subs.items()
-                    if var in rddl.state_fluents}
+                    if var in observed_vars}
             flat_subs = jax.tree_map(jnp.ravel, subs)
             states = list(flat_subs.values())
             state = jnp.concatenate(states)
@@ -841,6 +932,10 @@ class JaxDeepReactivePolicy(JaxPlan):
         def _jax_wrapped_drp_predict_train(key, params, hyperparams, step, subs):
             state = _jax_wrapped_subs_to_state(subs)
             actions = predict_fn.apply(params, state)
+            if not wrap_non_bool:
+                for (var, action) in actions.items():
+                    if var != bool_key and ranges[var] != 'bool':
+                        actions[var] = jnp.clip(action, *bounds[var])
             if use_constraint_satisfaction:
                 bool_actions = _jax_unstack_bool_from_softmax(actions[bool_key])
                 actions.update(bool_actions)
@@ -886,14 +981,14 @@ class JaxDeepReactivePolicy(JaxPlan):
         def _jax_wrapped_drp_init(key, hyperparams, subs):
             subs = {var: value[0, ...] 
                     for (var, value) in subs.items()
-                    if var in rddl.state_fluents}
+                    if var in observed_vars}
             state = _jax_wrapped_subs_to_state(subs)
             params = predict_fn.init(key, state)
             return params
         
         self.initializer = _jax_wrapped_drp_init
         
-    def guess_next_epoch(self, params: Dict) -> Dict:
+    def guess_next_epoch(self, params: Pytree) -> Pytree:
         return params
 
     
@@ -904,24 +999,135 @@ class JaxDeepReactivePolicy(JaxPlan):
 # - more stable but slower line search based planner
 #
 # ***********************************************************************
+
+class RollingMean:
+    '''Maintains an estimate of the rolling mean of a stream of real-valued 
+    observations.'''
+    
+    def __init__(self, window_size: int) -> None:
+        self._window_size = window_size
+        self._memory = deque(maxlen=window_size)
+        self._total = 0
+    
+    def update(self, x: float) -> float:
+        memory = self._memory
+        self._total += x
+        if len(memory) == self._window_size:
+            self._total -= memory.popleft()
+        memory.append(x)
+        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) -> None:
+        self._fig, axes = plt.subplots(1 + len(rddl.action_fluents))
+        
+        # prepare the loss plot
+        self._loss_ax = axes[0]
+        self._loss_ax.autoscale(enable=True)
+        self._loss_ax.set_xlabel('decision epoch')
+        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 action plots
+        self._action_ax = {name: axes[idx + 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()}
+        
+        plt.tight_layout()
+        plt.show(block=False)
+        
+    def redraw(self, xticks, losses, actions) -> 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 actions
+        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._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, 
+    whether progress was made, budget was reached, or other information that
+    can be used to monitor and act based on the planner's progress.'''
+    
+    NORMAL = 0
+    NO_PROGRESS = 1
+    PRECONDITION_POSSIBLY_UNSATISFIED = 2
+    TIME_BUDGET_REACHED = 3
+    ITER_BUDGET_REACHED = 4
+    INVALID_GRADIENT = 5
+    
+    def is_failure(self) -> bool:
+        return self.value >= 3
+
+
 class JaxBackpropPlanner:
     '''A class for optimizing an action sequence in the given RDDL MDP using 
     gradient descent.'''
     
     def __init__(self, rddl: RDDLLiftedModel,
                  plan: JaxPlan,
-                 batch_size_train: int,
-                 batch_size_test: int=None,
-                 rollout_horizon: int=None,
+                 batch_size_train: int=32,
+                 batch_size_test: Optional[int]=None,
+                 rollout_horizon: Optional[int]=None,
                  use64bit: bool=False,
-                 action_bounds: Dict[str, Tuple[np.ndarray, np.ndarray]]={},
+                 action_bounds: Optional[Bounds]=None,
                  optimizer: Callable[..., optax.GradientTransformation]=optax.rmsprop,
-                 optimizer_kwargs: Dict[str, object]={'learning_rate': 0.1},
-                 clip_grad: float=None,
+                 optimizer_kwargs: Optional[Kwargs]=None,
+                 clip_grad: Optional[float]=None,
                  logic: FuzzyLogic=FuzzyLogic(),
                  use_symlog_reward: bool=False,
-                 utility=jnp.mean,
-                 cpfs_without_grad: Set=set()) -> None:
+                 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:
         '''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
@@ -946,9 +1152,16 @@ class JaxBackpropPlanner:
         :param use_symlog_reward: whether to use the symlog transform on the 
         reward as a form of normalization
         :param utility: how to aggregate return observations to compute utility
-        of a policy or plan
+        of a policy or plan; must be either a function mapping jax array to a 
+        scalar, or a a string identifying the utility function by name 
+        ("mean", "mean_var", "entropic", or "cvar" are currently supported)
+        :param utility_kwargs: additional keyword arguments to pass hyper-
+        parameters to the utility function call
         :param cpfs_without_grad: which CPFs do not have gradients (use straight
         through gradient trick)
+        :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
         '''
         self.rddl = rddl
         self.plan = plan
@@ -959,22 +1172,25 @@ class JaxBackpropPlanner:
         if rollout_horizon is None:
             rollout_horizon = rddl.horizon
         self.horizon = rollout_horizon
+        if action_bounds is None:
+            action_bounds = {}
         self._action_bounds = action_bounds
         self.use64bit = use64bit
         self._optimizer_name = optimizer
+        if optimizer_kwargs is None:
+            optimizer_kwargs = {'learning_rate': 0.1}
         self._optimizer_kwargs = optimizer_kwargs
         self.clip_grad = clip_grad
         
         # set optimizer
         try:
             optimizer = optax.inject_hyperparams(optimizer)(**optimizer_kwargs)
-        except:
+        except Exception as _:
             raise_warning(
                 'Failed to inject hyperparameters into optax optimizer, '
                 'rolling back to safer method: please note that modification of '
                 'optimizer hyperparameters will not work, and it is '
-                'recommended to update your packages and Python distribution.',
-                'red')
+                'recommended to update optax and related packages.', 'red')
             optimizer = optimizer(**optimizer_kwargs)     
         if clip_grad is None:
             self.optimizer = optimizer
@@ -983,22 +1199,68 @@ class JaxBackpropPlanner:
                 optax.clip(clip_grad),
                 optimizer
             )
-            
+        
+        # set utility
+        if isinstance(utility, str):
+            utility = utility.lower()
+            if utility == 'mean':
+                utility_fn = jnp.mean
+            elif utility == 'mean_var':
+                utility_fn = mean_variance_utility
+            elif utility == 'entropic':
+                utility_fn = entropic_utility
+            elif utility == 'cvar':
+                utility_fn = cvar_utility
+            else:
+                raise RDDLNotImplementedError(
+                    f'Utility function <{utility}> is not supported: '
+                    'must be one of ["mean", "mean_var", "entropic", "cvar"].')
+        else:
+            utility_fn = utility
+        self.utility = utility_fn
+        
+        if utility_kwargs is None:
+            utility_kwargs = {}
+        self.utility_kwargs = utility_kwargs    
+        
         self.logic = logic
+        self.logic.set_use64bit(self.use64bit)
         self.use_symlog_reward = use_symlog_reward
-        self.utility = utility
+        if cpfs_without_grad is None:
+            cpfs_without_grad = set()
         self.cpfs_without_grad = cpfs_without_grad
+        self.compile_non_fluent_exact = compile_non_fluent_exact
+        self.logger = logger
         
         self._jax_compile_rddl()        
         self._jax_compile_optimizer()
-        
-    def summarize_hyperparameters(self):
-        print(f'objective and relaxations:\n'
-              f'    objective_fn    ={self.utility.__name__}\n'
+    
+    def _summarize_system(self) -> None:
+        try:
+            jaxlib_version = jax._src.lib.version_str
+        except Exception as _:
+            jaxlib_version = 'N/A'
+        try:
+            devices_short = ', '.join(
+                map(str, jax._src.xla_bridge.devices())).replace('\n', '')
+        except Exception as _:
+            devices_short = 'N/A'
+        print('\n'
+              f'JAX Planner version {__version__}\n' 
+              f'Python {sys.version}\n'
+              f'jax {jax.version.__version__}, jaxlib {jaxlib_version}, '
+              f'numpy {np.__version__}\n'
+              f'devices: {devices_short}\n')
+        
+    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'    model relaxation={type(self.logic).__name__}\n'
               f'    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'
@@ -1010,6 +1272,10 @@ class JaxBackpropPlanner:
         self.plan.summarize_hyperparameters()
         self.logic.summarize_hyperparameters()
         
+    # ===========================================================================
+    # COMPILATION SUBROUTINES
+    # ===========================================================================
+
     def _jax_compile_rddl(self):
         rddl = self.rddl
         
@@ -1017,13 +1283,18 @@ class JaxBackpropPlanner:
         self.compiled = JaxRDDLCompilerWithGrad(
             rddl=rddl,
             logic=self.logic,
+            logger=self.logger,
             use64bit=self.use64bit,
-            cpfs_without_grad=self.cpfs_without_grad)
-        self.compiled.compile()
+            cpfs_without_grad=self.cpfs_without_grad,
+            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, use64bit=self.use64bit)
-        self.test_compiled.compile()
+        self.test_compiled = JaxRDDLCompiler(
+            rddl=rddl, 
+            logger=self.logger,
+            use64bit=self.use64bit)
+        self.test_compiled.compile(log_jax_expr=True, heading='EXACT MODEL')
         
     def _jax_compile_optimizer(self):
         
@@ -1051,11 +1322,10 @@ class JaxBackpropPlanner:
         
         # losses
         train_loss = self._jax_loss(train_rollouts, use_symlog=self.use_symlog_reward)
-        self.train_loss = jax.jit(train_loss)
         self.test_loss = jax.jit(self._jax_loss(test_rollouts, use_symlog=False))
         
         # optimization
-        self.update = jax.jit(self._jax_update(train_loss))
+        self.update = self._jax_update(train_loss)
     
     def _jax_return(self, use_symlog):
         gamma = self.rddl.discount
@@ -1068,13 +1338,14 @@ class JaxBackpropPlanner:
                 rewards = rewards * discount[jnp.newaxis, ...]
             returns = jnp.sum(rewards, axis=1)
             if use_symlog:
-                returns = jnp.sign(returns) * jnp.log1p(jnp.abs(returns))
+                returns = jnp.sign(returns) * jnp.log(1.0 + jnp.abs(returns))
             return returns
         
         return _jax_wrapped_returns
         
     def _jax_loss(self, rollouts, use_symlog=False): 
-        utility_fn = self.utility        
+        utility_fn = self.utility    
+        utility_kwargs = self.utility_kwargs 
         _jax_wrapped_returns = self._jax_return(use_symlog)
         
         # the loss is the average cumulative reward across all roll-outs
@@ -1083,7 +1354,7 @@ class JaxBackpropPlanner:
             log = rollouts(key, policy_params, hyperparams, subs, model_params)
             rewards = log['reward']
             returns = _jax_wrapped_returns(rewards)
-            utility = utility_fn(returns)
+            utility = utility_fn(returns, **utility_kwargs)
             loss = -utility
             return loss, log
         
@@ -1096,7 +1367,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
+            return policy_params, opt_state, None
         
         return _jax_wrapped_init_policy
         
@@ -1107,17 +1378,18 @@ class JaxBackpropPlanner:
         # 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,
-                                     subs, model_params, opt_state):
-            grad_fn = jax.grad(loss, argnums=1, has_aux=True)
-            grad, log = grad_fn(key, policy_params, hyperparams, subs, model_params)  
+                                     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) 
             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, log
+            return policy_params, converged, opt_state, None, loss_val, log
         
-        return _jax_wrapped_plan_update
+        return jax.jit(_jax_wrapped_plan_update)
             
     def _batched_init_subs(self, subs): 
         rddl = self.rddl
@@ -1145,13 +1417,15 @@ class JaxBackpropPlanner:
         
         return init_train, init_test
     
-    def optimize(self, *args, return_callback: bool=False, **kwargs) -> object:
-        ''' Compute an optimal straight-line plan. Returns the parameters
-        for the optimized policy.
+    # ===========================================================================
+    # OPTIMIZE API
+    # ===========================================================================
+
+    def optimize(self, *args, **kwargs) -> Dict[str, Any]:
+        ''' Compute an optimal policy or plan. Return the callback from training.
         
-        :param key: JAX PRNG key   
+        :param key: JAX PRNG key (derived from clock if not provided)
         :param epochs: the maximum number of steps of gradient descent
-        :param 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 model_params: optional model-parameters to override default
@@ -1162,33 +1436,44 @@ class JaxBackpropPlanner:
         :param guess: initial policy parameters: if None will use the initializer
         specified in this instance
         :param verbose: not print (0), print summary (1), print progress (2)
-        :param return_callback: whether to return the callback from training
-        instead of the parameters
+        :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
+        :param tqdm_position: position of tqdm progress bar (for multiprocessing)
         '''
         it = self.optimize_generator(*args, **kwargs)
-        callback = deque(it, maxlen=1).pop()
-        if return_callback:
-            return callback
+        
+        # if the python is C-compiled then the deque is native C and much faster
+        # than naively exhausting iterator, but not if the python is some other
+        # version (e.g. PyPi); for details, see
+        # https://stackoverflow.com/questions/50937966/fastest-most-pythonic-way-to-consume-an-iterator
+        callback = None
+        if sys.implementation.name == 'cpython':
+            last_callback = deque(it, maxlen=1)
+            if last_callback:
+                callback = last_callback.pop()
         else:
-            return callback['best_params']
+            for callback in it:
+                pass
+        return callback
     
-    def optimize_generator(self, key: random.PRNGKey,
+    def optimize_generator(self, key: Optional[random.PRNGKey]=None,
                            epochs: int=999999,
                            train_seconds: float=120.,
-                           plot_step: int=None,
-                           model_params: Dict[str, object]=None,
-                           policy_hyperparams: Dict[str, object]=None,
-                           subs: Dict[str, object]=None,
-                           guess: Dict[str, object]=None,
+                           plot_step: Optional[int]=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,
                            verbose: int=2,
-                           tqdm_position: int=None) -> Generator[Dict[str, object], None, None]:
-        '''Returns a generator for computing an optimal straight-line plan. 
+                           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. 
         Generator can be iterated over to lazily optimize the plan, yielding
         a dictionary of intermediate computations.
         
-        :param key: JAX PRNG key   
+        :param key: JAX PRNG key (derived from clock if not provided)
         :param epochs: the maximum number of steps of gradient descent
-        :param 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 model_params: optional model-parameters to override default
@@ -1199,26 +1484,53 @@ class JaxBackpropPlanner:
         :param guess: initial policy parameters: if None will use the initializer
         specified in this instance
         :param verbose: not print (0), print summary (1), print progress (2)
+        :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
         :param tqdm_position: position of tqdm progress bar (for multiprocessing)
         '''
         verbose = int(verbose)
         start_time = time.time()
         elapsed_outside_loop = 0
         
+        # if PRNG key is not provided
+        if key is None:
+            key = random.PRNGKey(round(time.time() * 1000))
+            
+        # if policy_hyperparams is not provided
+        if policy_hyperparams is None:
+            raise_warning('policy_hyperparams is not set, setting 1.0 for '
+                          'all action-fluents which could be suboptimal.')
+            policy_hyperparams = {action: 1.0 
+                                  for action in self.rddl.action_fluents}
+        
+        # if policy_hyperparams is a scalar
+        elif isinstance(policy_hyperparams, (int, float, np.number)):
+            raise_warning(f'policy_hyperparams is {policy_hyperparams}, '
+                          'setting this value for all action-fluents.')
+            hyperparam_value = float(policy_hyperparams)
+            policy_hyperparams = {action: hyperparam_value
+                                  for action in self.rddl.action_fluents}
+            
         # print summary of parameters:
         if verbose >= 1:
-            print('==============================================\n'
-                  'JAX PLANNER PARAMETER SUMMARY\n'
-                  '==============================================')
+            self._summarize_system()
             self.summarize_hyperparameters()
             print(f'optimize() call hyper-parameters:\n'
+                  f'    PRNG key           ={key}\n'
                   f'    max_iterations     ={epochs}\n'
                   f'    max_seconds        ={train_seconds}\n'
                   f'    model_params       ={model_params}\n'
                   f'    policy_hyper_params={policy_hyperparams}\n'
                   f'    override_subs_dict ={subs is not None}\n'
-                  f'    provide_param_guess={guess is not None}\n' 
-                  f'    plot_frequency     ={plot_step}\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'    verbose            ={verbose}\n')
+            if verbose >= 2 and self.compiled.relaxations:
+                print('Some RDDL operations are non-differentiable, '
+                      'replacing them with differentiable relaxations:')
+                print(self.compiled.summarize_model_relaxations())
             
         # compute a batched version of the initial values
         if subs is None:
@@ -1245,14 +1557,26 @@ class JaxBackpropPlanner:
         # initialize policy parameters
         if guess is None:
             key, subkey = random.split(key)
-            policy_params, opt_state = self.initialize(
+            policy_params, opt_state, opt_aux = self.initialize(
                 subkey, policy_hyperparams, train_subs)
         else:
             policy_params = guess
             opt_state = self.optimizer.init(policy_params)
+            opt_aux = None
+        
+        # initialize running statistics
         best_params, best_loss, best_grad = policy_params, jnp.inf, jnp.inf
         last_iter_improve = 0
+        rolling_test_loss = RollingMean(test_rolling_window)
         log = {}
+        status = JaxPlannerStatus.NORMAL
+        
+        # initialize plot area
+        if plot_step is None or plot_step <= 0 or plt is None:
+            plot = None
+        else:
+            plot = JaxPlannerPlot(self.rddl, self.horizon)
+        xticks, loss_values = [], []
         
         # training loop
         iters = range(epochs)
@@ -1260,25 +1584,25 @@ class JaxBackpropPlanner:
             iters = tqdm(iters, total=100, position=tqdm_position)
         
         for it in iters:
+            status = JaxPlannerStatus.NORMAL
             
             # update the parameters of the plan
-            key, subkey1, subkey2, subkey3 = random.split(key, num=4)
-            policy_params, converged, opt_state, train_log = self.update(
-                subkey1, policy_params, policy_hyperparams,
-                train_subs, model_params, opt_state)
+            key, subkey = random.split(key)
+            policy_params, converged, opt_state, opt_aux, train_loss, train_log = \
+                self.update(subkey, policy_params, policy_hyperparams,
+                            train_subs, model_params, opt_state, opt_aux)
             if not np.all(converged):
                 raise_warning(
                     'Projected gradient method for satisfying action concurrency '
                     'constraints reached the iteration limit: plan is possibly '
                     'invalid for the current instance.', 'red')
+                status = JaxPlannerStatus.PRECONDITION_POSSIBLY_UNSATISFIED
             
             # evaluate losses
-            train_loss, _ = self.train_loss(
-                subkey2, policy_params, policy_hyperparams,
-                train_subs, model_params)
             test_loss, log = self.test_loss(
-                subkey3, policy_params, policy_hyperparams,
+                subkey, policy_params, policy_hyperparams,
                 test_subs, model_params_test)
+            test_loss = rolling_test_loss.update(test_loss)
             
             # record the best plan so far
             if test_loss < best_loss:
@@ -1287,21 +1611,45 @@ class JaxBackpropPlanner:
                 last_iter_improve = it
             
             # save the plan figure
-            if plot_step is not None and it % plot_step == 0:
-                self._plot_actions(
-                    key, policy_params, policy_hyperparams, test_subs, it)
+            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}
+                plot.redraw(xticks, loss_values, action_values)
             
             # if the progress bar is used
             elapsed = time.time() - start_time - elapsed_outside_loop
             if verbose >= 2:
                 iters.n = int(100 * min(1, max(elapsed / train_seconds, it / epochs)))
                 iters.set_description(
-                    f'[{tqdm_position}] {it:6} it / {-train_loss:14.4f} train / '
-                    f'{-test_loss:14.4f} test / {-best_loss:14.4f} best')
+                    f'[{tqdm_position}] {it:6} it / {-train_loss:14.6f} train / '
+                    f'{-test_loss:14.6f} test / {-best_loss:14.6f} best')
+                        
+            # reached computation budget
+            if elapsed >= train_seconds:
+                status = JaxPlannerStatus.TIME_BUDGET_REACHED
+            if it >= epochs - 1:
+                status = JaxPlannerStatus.ITER_BUDGET_REACHED
+            
+            # numerical error
+            if not np.isfinite(train_loss):
+                raise_warning(
+                    f'Aborting JAX planner due to invalid train loss {train_loss}.',
+                    'red')
+                status = JaxPlannerStatus.INVALID_GRADIENT
+            
+            # no progress
+            grad_norm_zero, _ = jax.tree_util.tree_flatten(
+                jax.tree_map(lambda x: np.allclose(x, 0), train_log['grad']))
+            if np.all(grad_norm_zero):
+                status = JaxPlannerStatus.NO_PROGRESS
             
             # return a callback
             start_time_outside = time.time()
             yield {
+                'status': status,
                 'iteration': it,
                 'train_return':-train_loss,
                 'test_return':-test_loss,
@@ -1318,16 +1666,15 @@ class JaxBackpropPlanner:
             }
             elapsed_outside_loop += (time.time() - start_time_outside)
             
-            # reached time budget
-            if elapsed >= train_seconds:
-                break
-            
-            # numerical error
-            if not np.isfinite(train_loss):
+            # abortion check
+            if status.is_failure():
                 break
-        
+                        
+        # release resources
         if verbose >= 2:
             iters.close()
+        if plot is not None:
+            plot.close()
         
         # validate the test return
         if log:
@@ -1337,24 +1684,23 @@ class JaxBackpropPlanner:
             if messages:
                 messages = '\n'.join(messages)
                 raise_warning('The JAX compiler encountered the following '
-                              'problems in the original RDDL '
+                              'error(s) in the original RDDL formulation '
                               f'during test evaluation:\n{messages}', 'red')                               
         
         # summarize and test for convergence
         if verbose >= 1:
-            grad_norm = jax.tree_map(
-                lambda x: np.array(jnp.linalg.norm(x)).item(), best_grad)
+            grad_norm = jax.tree_map(lambda x: np.linalg.norm(x).item(), best_grad)
             diagnosis = self._perform_diagnosis(
-                last_iter_improve, it, 
-                -train_loss, -test_loss, -best_loss, grad_norm)
+                last_iter_improve, -train_loss, -test_loss, -best_loss, grad_norm)
             print(f'summary of optimization:\n'
+                  f'    status_code   ={status}\n'
                   f'    time_elapsed  ={elapsed}\n'
                   f'    iterations    ={it}\n'
                   f'    best_objective={-best_loss}\n'
-                  f'    grad_norm     ={grad_norm}\n'
+                  f'    best_grad_norm={grad_norm}\n'
                   f'diagnosis: {diagnosis}\n')
     
-    def _perform_diagnosis(self, last_iter_improve, total_it, 
+    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)
@@ -1373,20 +1719,20 @@ class JaxBackpropPlanner:
             if grad_is_zero:
                 return termcolor.colored(
                     '[FAILURE] no progress was made, '
-                    f'and max grad norm = {max_grad_norm}, '
-                    'likely stuck in a plateau.', 'red')
+                    f'and max grad norm {max_grad_norm:.6f} is zero: '
+                    'solver likely stuck in a plateau.', 'red')
             else:
                 return termcolor.colored(
                     '[FAILURE] no progress was made, '
-                    f'but max grad norm = {max_grad_norm} > 0, '
-                    'likely due to bad l.r. or other hyper-parameter.', 'red')
+                    f'but max grad norm {max_grad_norm:.6f} is non-zero: '
+                    'likely poor learning rate or other hyper-parameter.', 'red')
         
         # model is likely poor IF:
         # 1. the train and test return disagree
         if not (validation_error < 20):
             return termcolor.colored(
                 '[WARNING] progress was made, '
-                f'but relative train test error = {validation_error} is high, '
+                f'but relative train-test error {validation_error:.6f} is high: '
                 'likely poor model relaxation around the solution, '
                 'or the batch size is too small.', 'yellow')
         
@@ -1397,208 +1743,216 @@ class JaxBackpropPlanner:
             if not (return_to_grad_norm > 1):
                 return termcolor.colored(
                     '[WARNING] progress was made, '
-                    f'but max grad norm = {max_grad_norm} is high, '
-                    'likely indicates the solution is not locally optimal, '
-                    'or the model is not smooth around the solution, '
+                    f'but max grad norm {max_grad_norm:.6f} is high: '
+                    'likely the solution is not locally optimal, '
+                    'or the relaxed model is not smooth around the solution, '
                     'or the batch size is too small.', 'yellow')
         
         # likely successful
         return termcolor.colored(
-            '[SUCCESS] planner appears to have converged successfully '
+            '[SUCCESS] planner has converged successfully '
             '(note: not all potential problems can be ruled out).', 'green')
         
     def get_action(self, key: random.PRNGKey,
-                   params: Dict,
+                   params: Pytree,
                    step: int,
-                   subs: Dict,
-                   policy_hyperparams: Dict[str, object]=None) -> Dict[str, object]:
+                   subs: Dict[str, Any],
+                   policy_hyperparams: Optional[Dict[str, Any]]=None) -> Dict[str, Any]:
         '''Returns an action dictionary from the policy or plan with the given
         parameters.
         
         :param key: the JAX PRNG key
         :param params: the trainable parameter PyTree of the policy
         :param step: the time step at which decision is made
-        :param policy_hyperparams: hyper-parameters for the policy/plan, such as
-        weights for sigmoid wrapping boolean actions
         :param subs: the dict of pvariables
+        :param policy_hyperparams: hyper-parameters for the policy/plan, such as
+        weights for sigmoid wrapping boolean actions (optional)
         '''
         
         # check compatibility of the subs dictionary
-        for var in subs.keys():
+        for (var, values) in subs.items():
+            
+            # must not be grounded
             if RDDLPlanningModel.FLUENT_SEP in var \
             or RDDLPlanningModel.OBJECT_SEP in var:
-                raise Exception(f'State dictionary passed to the JAX policy is '
-                                f'grounded, since it contains the key <{var}>, '
-                                f'but a vectorized environment is required: '
-                                f'please make sure vectorized=True in the RDDLEnv.')
-        
+                raise ValueError(f'State dictionary passed to the JAX policy is '
+                                 f'grounded, since it contains the key <{var}>, '
+                                 f'but a vectorized environment is required: '
+                                 f'please make sure vectorized=True in the RDDLEnv.')
+            
+            # must be numeric array
+            # exception is for POMDPs at 1st epoch when observ-fluents are None
+            if not jnp.issubdtype(values.dtype, jnp.number) \
+            and not jnp.issubdtype(values.dtype, jnp.bool_):
+                if step == 0 and var in self.rddl.observ_fluents:
+                    subs[var] = self.test_compiled.init_values[var]
+                else:
+                    raise ValueError(f'Values assigned to pvariable {var} are '
+                                     f'non-numeric of type {values.dtype}: {values}.')
+            
         # cast device arrays to numpy
         actions = self.test_policy(key, params, policy_hyperparams, step, subs)
         actions = jax.tree_map(np.asarray, actions)
         return actions      
-            
-    def _plot_actions(self, key, params, hyperparams, subs, it):
-        rddl = self.rddl
-        try:
-            import matplotlib.pyplot as plt
-        except Exception:
-            print('matplotlib is not installed, aborting plot...')
-            return
-            
-        # predict actions from the trained policy or plan
-        actions = self.test_rollouts(key, params, hyperparams, subs, {})['action']
-            
-        # plot the action sequences as color maps
-        fig, axs = plt.subplots(nrows=len(actions), constrained_layout=True)
-        for (ax, name) in zip(axs, actions):
-            action = np.mean(actions[name], axis=0, dtype=float)
-            action = np.reshape(action, newshape=(action.shape[0], -1)).T
-            if rddl.variable_ranges[name] == 'bool':
-                vmin, vmax = 0.0, 1.0
-            else:
-                vmin, vmax = None, None                
-            img = ax.imshow(
-                action, vmin=vmin, vmax=vmax, cmap='seismic', aspect='auto')
-            ax.set_xlabel('time')
-            ax.set_ylabel(name)
-            plt.colorbar(img, ax=ax)
-            
-        # write plot to disk
-        plt.savefig(f'plan_{rddl.domain_name}_{rddl.instance_name}_{it}.pdf',
-                    bbox_inches='tight')
-        plt.clf()
-        plt.close(fig)
     
 
-class JaxArmijoLineSearchPlanner(JaxBackpropPlanner):
+class JaxLineSearchPlanner(JaxBackpropPlanner):
     '''A class for optimizing an action sequence in the given RDDL MDP using 
-    Armijo linear search gradient descent.'''
+    linear search gradient descent, with the Armijo condition.'''
     
     def __init__(self, *args,
                  optimizer: Callable[..., optax.GradientTransformation]=optax.sgd,
-                 optimizer_kwargs: Dict[str, object]={'learning_rate': 1.0},
-                 beta: float=0.8,
+                 optimizer_kwargs: Kwargs={'learning_rate': 1.0},
+                 decay: float=0.8,
                  c: float=0.1,
-                 lrmax: float=1.0,
-                 lrmin: float=1e-5,
+                 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 Armijo line search. All arguments are the
+        (plan) in the given RDDL using line search. All arguments are the
         same as in the parent class, except:
         
-        :param beta: reduction factor of learning rate per line search iteration
-        :param c: coefficient in Armijo condition
-        :param lrmax: initial learning rate for line search
-        :param lrmin: minimum possible learning rate (line search halts)
+        :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.beta = beta
+        self.decay = decay
         self.c = c
-        self.lrmax = lrmax
-        self.lrmin = lrmin
-        super(JaxArmijoLineSearchPlanner, self).__init__(
+        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,
             optimizer=optimizer,
             optimizer_kwargs=optimizer_kwargs,
             **kwargs)
         
-    def summarize_hyperparameters(self):
-        super(JaxArmijoLineSearchPlanner, self).summarize_hyperparameters()
+    def summarize_hyperparameters(self) -> None:
+        super(JaxLineSearchPlanner, self).summarize_hyperparameters()
         print(f'linesearch hyper-parameters:\n'
-              f'    beta    ={self.beta}\n'
+              f'    decay   ={self.decay}\n'
               f'    c       ={self.c}\n'
-              f'    lr_range=({self.lrmin}, {self.lrmax})\n')
+              f'    lr_range=({self.step_min}, {self.step_max})')
     
     def _jax_update(self, loss):
         optimizer = self.optimizer
         projection = self.plan.projection
-        beta, c, lrmax, lrmin = self.beta, self.c, self.lrmax, self.lrmin
-        
-        # continue line search if Armijo condition not satisfied and learning
-        # rate can be further reduced
-        def _jax_wrapped_line_search_armijo_check(val):
-            (_, old_f, _, old_norm_g2, _), (_, new_f, lr, _), _, _ = val            
-            return jnp.logical_and(
-                new_f >= old_f - c * lr * old_norm_g2,
-                lr >= lrmin / beta)
-            
-        def _jax_wrapped_line_search_iteration(val):
-            old, new, best, aux = val
-            old_x, _, old_g, _, old_state = old
-            _, _, lr, iters = new
-            _, best_f, _, _ = best
-            key, hyperparams, *other = aux
-            
-            # anneal learning rate and apply a gradient step
-            new_lr = beta * lr
-            old_state.hyperparams['learning_rate'] = new_lr
-            updates, new_state = optimizer.update(old_g, old_state)
-            new_x = optax.apply_updates(old_x, updates)
-            new_x, _ = projection(new_x, hyperparams)
-            
-            # evaluate new loss and record best so far
-            new_f, _ = loss(key, new_x, hyperparams, *other)
-            new = (new_x, new_f, new_lr, iters + 1)
-            best = jax.lax.cond(
-                new_f < best_f,
-                lambda: (new_x, new_f, new_lr, new_state),
-                lambda: best
-            )
-            return old, new, best, aux
+        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):
-            
-            # calculate initial loss value, gradient and squared norm
-            old_x = policy_params
-            loss_and_grad_fn = jax.value_and_grad(loss, argnums=1, has_aux=True)
-            (old_f, log), old_g = loss_and_grad_fn(
-                key, old_x, hyperparams, subs, model_params)            
-            old_norm_g2 = jax.tree_map(lambda x: jnp.sum(jnp.square(x)), old_g)
-            old_norm_g2 = jax.tree_util.tree_reduce(jnp.add, old_norm_g2)
-            log['grad'] = old_g
+                                     subs, model_params, opt_state, opt_aux):
             
-            # initialize learning rate to maximum
-            new_lr = lrmax / beta
-            old = (old_x, old_f, old_g, old_norm_g2, opt_state)
-            new = (old_x, old_f, new_lr, 0)            
-            best = (old_x, jnp.inf, jnp.nan, opt_state)
-            aux = (key, hyperparams, subs, model_params)
+            # initialize the line search
+            f, grad, gnorm2, log = _jax_wrapped_line_search_init(
+                key, policy_params, hyperparams, subs, model_params)            
             
-            # do a single line search step with the initial learning rate
-            init_val = (old, new, best, aux)            
-            init_val = _jax_wrapped_line_search_iteration(init_val)
+            # 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:
+                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
             
-            # continue to anneal the learning rate until Armijo condition holds
-            # or the learning rate becomes too small, then use the best parameter
-            _, (*_, iters), (best_params, _, best_lr, best_state), _ = \
-            jax.lax.while_loop(
-                cond_fun=_jax_wrapped_line_search_armijo_check,
-                body_fun=_jax_wrapped_line_search_iteration,
-                init_val=init_val
-            )
-            best_state.hyperparams['learning_rate'] = best_lr
             log['updates'] = None
-            log['line_search_iters'] = iters
-            log['learning_rate'] = best_lr
-            return best_params, True, best_state, log
+            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
+# 
+# Based on the original paper "A Distributional Framework for Risk-Sensitive 
+# End-to-End Planning in Continuous MDPs" by Patton et al., AAAI 2022.
+#
+# Original risk functions:
+# - entropic utility
+# - mean-variance approximation
+# - conditional value at risk with straight-through gradient trick
+#
+# ***********************************************************************
+
+
+@jax.jit
+def entropic_utility(returns: jnp.ndarray, beta: float) -> float:
+    return (-1.0 / beta) * jax.scipy.special.logsumexp(
+            -beta * returns, b=1.0 / returns.size)
+
+
+@jax.jit
+def mean_variance_utility(returns: jnp.ndarray, beta: float) -> float:
+    return jnp.mean(returns) - (beta / 2.0) * jnp.var(returns)
+    
+
+@jax.jit
+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
+# 
+# - offline controller is the straight-line planner
+# - online controller is the replanning mode
+#
+# ***********************************************************************
+
 class JaxOfflineController(BaseAgent):
     '''A container class for a Jax policy trained offline.'''
+    
     use_tensor_obs = True
     
-    def __init__(self, planner: JaxBackpropPlanner, key: random.PRNGKey,
-                 eval_hyperparams: Dict[str, object]=None,
-                 params: Dict[str, object]=None,
+    def __init__(self, planner: JaxBackpropPlanner, 
+                 key: Optional[random.PRNGKey]=None,
+                 eval_hyperparams: Optional[Dict[str, Any]]=None,
+                 params: Optional[Pytree]=None,
                  train_on_reset: bool=False,
                  **train_kwargs) -> None:
         '''Creates a new JAX offline control policy that is trained once, then
         deployed later.
         
         :param planner: underlying planning algorithm for optimizing actions
-        :param key: the RNG key to seed randomness
+        :param key: the RNG key to seed randomness (derives from clock if not
+        provided)
         :param eval_hyperparams: policy hyperparameters to apply for evaluation
         or whenever sample_action is called
         :param params: use the specified policy parameters instead of calling
@@ -1608,6 +1962,8 @@ class JaxOfflineController(BaseAgent):
         for optimization
         '''
         self.planner = planner
+        if key is None:
+            key = random.PRNGKey(round(time.time() * 1000))
         self.key = key
         self.eval_hyperparams = eval_hyperparams
         self.train_on_reset = train_on_reset
@@ -1616,17 +1972,18 @@ class JaxOfflineController(BaseAgent):
         
         self.step = 0
         if not self.train_on_reset and not self.params_given:
-            params = self.planner.optimize(key=self.key, **self.train_kwargs) 
+            callback = self.planner.optimize(key=self.key, **self.train_kwargs)
+            params = callback['best_params'] 
         self.params = params  
         
-    def sample_action(self, state):
+    def sample_action(self, state: Dict[str, Any]) -> Dict[str, Any]:
         self.key, subkey = random.split(self.key)
         actions = self.planner.get_action(
             subkey, self.params, self.step, state, self.eval_hyperparams)
         self.step += 1
         return actions
         
-    def reset(self):
+    def reset(self) -> None:
         self.step = 0
         if self.train_on_reset and not self.params_given:
             self.params = self.planner.optimize(key=self.key, **self.train_kwargs)
@@ -1635,41 +1992,51 @@ class JaxOfflineController(BaseAgent):
 class JaxOnlineController(BaseAgent):
     '''A container class for a Jax controller continuously updated using state 
     feedback.'''
+    
     use_tensor_obs = True
     
-    def __init__(self, planner: JaxBackpropPlanner, key: random.PRNGKey,
-                 eval_hyperparams: Dict=None, warm_start: bool=True,
+    def __init__(self, planner: JaxBackpropPlanner, 
+                 key: Optional[random.PRNGKey]=None,
+                 eval_hyperparams: Optional[Dict[str, Any]]=None,
+                 warm_start: bool=True,
                  **train_kwargs) -> None:
         '''Creates a new JAX control policy that is trained online in a closed-
         loop fashion.
         
         :param planner: underlying planning algorithm for optimizing actions
-        :param key: the RNG key to seed randomness
+        :param key: the RNG key to seed randomness (derives from clock if not
+        provided)
         :param eval_hyperparams: policy hyperparameters to apply for evaluation
         or whenever sample_action is called
+        :param warm_start: whether to use the previous decision epoch final
+        policy parameters to warm the next decision epoch
         :param **train_kwargs: any keyword arguments to be passed to the planner
         for optimization
         '''
         self.planner = planner
+        if key is None:
+            key = random.PRNGKey(round(time.time() * 1000))
         self.key = key
         self.eval_hyperparams = eval_hyperparams
         self.warm_start = warm_start
         self.train_kwargs = train_kwargs
         self.reset()
      
-    def sample_action(self, state):
+    def sample_action(self, state: Dict[str, Any]) -> Dict[str, Any]:
         planner = self.planner
-        params = planner.optimize(
+        callback = planner.optimize(
             key=self.key,
             guess=self.guess,
             subs=state,
             **self.train_kwargs)
+        params = callback['best_params']
         self.key, subkey = random.split(self.key)
-        actions = planner.get_action(subkey, params, 0, state, self.eval_hyperparams)
+        actions = planner.get_action(
+            subkey, params, 0, state, self.eval_hyperparams)
         if self.warm_start:
             self.guess = planner.plan.guess_next_epoch(params)
         return actions
         
-    def reset(self):
+    def reset(self) -> None:
         self.guess = None