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

pyRDDLGym_jax/core/planner.py

@@ -1,6 +1,7 @@
 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
@@ -12,12 +13,33 @@ import os
 import sys
 import termcolor
 import time
+import traceback
 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
+
+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
@@ -37,6 +59,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 +82,96 @@ 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, hk.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 +183,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 +206,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 
@@ -139,28 +221,37 @@ class JaxRDDLCompilerWithGrad(JaxRDDLCompiler):
         :param *kwargs: keyword arguments to pass to base compiler
         '''
         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.')   
+        pvars_cast = set()
         for (var, values) in self.init_values.items():
             self.init_values[var] = np.asarray(values, dtype=self.REAL) 
+            if not np.issubdtype(np.atleast_1d(values).dtype, np.floating):
+                pvars_cast.add(var)
+        if pvars_cast:
+            raise_warning(f'JAX gradient compiler requires that initial values '
+                          f'of p-variables {pvars_cast} be cast to float.')   
         
         # overwrite basic operations with fuzzy ones
         self.RELATIONAL_OPS = {
-            '>=': logic.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 +260,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):
@@ -188,46 +283,33 @@ class JaxRDDLCompilerWithGrad(JaxRDDLCompiler):
         return _jax_wrapped_stop_grad
         
     def _compile_cpfs(self, info):
-        raise_warning('CPFs outputs will be cast to real.')      
+        cpfs_cast = set()   
         jax_cpfs = {}
         for (_, cpfs) in self.levels.items():
             for cpf in cpfs:
                 _, expr = self.rddl.cpfs[cpf]
                 jax_cpfs[cpf] = self._jax(expr, info, dtype=self.REAL)
+                if self.rddl.variable_ranges[cpf] != 'real':
+                    cpfs_cast.add(cpf)
                 if cpf in self.cpfs_without_grad:
-                    raise_warning(f'CPF <{cpf}> stops gradient.')      
                     jax_cpfs[cpf] = self._jax_stop_grad(jax_cpfs[cpf])
+                    
+        if cpfs_cast:
+            raise_warning(f'JAX gradient compiler requires that outputs of CPFs '
+                          f'{cpfs_cast} be cast to float.') 
+        if self.cpfs_without_grad:
+            raise_warning(f'User requested that gradients not flow '
+                          f'through CPFs {self.cpfs_without_grad}.')      
         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.')            
-                       
+            raise_warning('JAX gradient compiler ignores KronDelta '
+                          'during compilation.')            
         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 +318,7 @@ class JaxRDDLCompilerWithGrad(JaxRDDLCompiler):
 # - deep reactive policy
 #
 # ***********************************************************************
+
 class JaxPlan:
     '''Base class for all JAX policy representations.'''
     
@@ -244,16 +327,17 @@ class JaxPlan:
         self._train_policy = None
         self._test_policy = None
         self._projection = None
-    
-    def summarize_hyperparameters(self):
+        self.bounds = None
+        
+    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 +373,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':
@@ -298,7 +383,7 @@ class JaxPlan:
             # check invalid type
             if prange not in compiled.JAX_TYPES:
                 raise RDDLTypeError(
-                    f'Invalid range <{prange}. of action-fluent <{name}>, '
+                    f'Invalid range <{prange}> of action-fluent <{name}>, '
                     f'must be one of {set(compiled.JAX_TYPES.keys())}.')
                 
             # clip boolean to (0, 1), otherwise use the RDDL action bounds
@@ -309,8 +394,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),
@@ -320,7 +405,7 @@ class JaxPlan:
                                     ~lower_finite & upper_finite,
                                     ~lower_finite & ~upper_finite]
             bounds[name] = (lower, upper)
-            raise_warning(f'Bounds of action fluent <{name}> set to {bounds[name]}.')
+            raise_warning(f'Bounds of action-fluent <{name}> set to {bounds[name]}.')
         return shapes, bounds, bounds_safe, cond_lists
     
     def _count_bool_actions(self, rddl: RDDLLiftedModel):
@@ -336,7 +421,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,
@@ -362,6 +447,7 @@ class JaxStraightLinePlan(JaxPlan):
         use_new_projection = True
         '''
         super(JaxStraightLinePlan, self).__init__()
+        
         self._initializer_base = initializer
         self._initializer = initializer
         self._wrap_sigmoid = wrap_sigmoid
@@ -371,10 +457,13 @@ 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:
+        bounds = '\n        '.join(
+            map(lambda kv: f'{kv[0]}: {kv[1]}', self.bounds.items()))
         print(f'policy hyper-parameters:\n'
-              f'    initializer          ={type(self._initializer_base).__name__}\n'
+              f'    initializer          ={self._initializer_base}\n'
               f'constraint-sat strategy (simple):\n'
+              f'    parsed_action_bounds =\n        {bounds}\n'
               f'    wrap_sigmoid         ={self._wrap_sigmoid}\n'
               f'    wrap_sigmoid_min_prob={self._min_action_prob}\n'
               f'    wrap_non_bool        ={self._wrap_non_bool}\n'
@@ -383,7 +472,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 +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.log1p(1.0 / action - 2.0)
+                return (-1.0 / weight) * jnp.log(1.0 / action - 1.0)
             else:
                 return action
             
@@ -506,7 +596,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)
@@ -537,7 +627,7 @@ class JaxStraightLinePlan(JaxPlan):
             if 1 < allowed_actions < bool_action_count:
                 raise RDDLNotImplementedError(
                     f'Straight-line plans with wrap_softmax currently '
-                    f'do not support max-nondef-actions = {allowed_actions} > 1.')
+                    f'do not support max-nondef-actions {allowed_actions} > 1.')
                 
             # potentially apply projection but to non-bool actions only
             self.projection = _jax_wrapped_slp_project_to_box
@@ -668,14 +758,14 @@ class JaxStraightLinePlan(JaxPlan):
             for (var, shape) in shapes.items():
                 if ranges[var] != 'bool' or not stack_bool_params: 
                     key, subkey = random.split(key)
-                    param = init(subkey, shape, dtype=compiled.REAL)
+                    param = init(key=subkey, shape=shape, dtype=compiled.REAL)
                     if ranges[var] == 'bool':
                         param += bool_threshold
                     params[var] = param
             if stack_bool_params:
                 key, subkey = random.split(key)
                 bool_shape = (horizon, bool_action_count)
-                bool_param = init(subkey, bool_shape, dtype=compiled.REAL)
+                bool_param = init(key=subkey, shape=bool_shape, dtype=compiled.REAL)
                 params[bool_key] = bool_param
             params, _ = _jax_wrapped_slp_project_to_box(params, hyperparams)
             return params
@@ -688,7 +778,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 +786,13 @@ 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=False, 
+                 normalize_per_layer: bool=False,
+                 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 +800,45 @@ 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 normalize_per_layer: whether to apply layer norm to each input
+        individually (only active if normalize is True)
+        :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
+        self._normalize_per_layer = normalize_per_layer
+        if normalizer_kwargs is None:
+            normalizer_kwargs = {'create_offset': True, 'create_scale': True}
+        self._normalizer_kwargs = normalizer_kwargs
+        self._wrap_non_bool = wrap_non_bool
             
-    def summarize_hyperparameters(self):
+    def summarize_hyperparameters(self) -> None:
+        bounds = '\n        '.join(
+            map(lambda kv: f'{kv[0]}: {kv[1]}', self.bounds.items()))
         print(f'policy hyper-parameters:\n'
-              f'    topology        ={self._topology}\n'
-              f'    activation_fn   ={self._activations[0].__name__}\n'
-              f'    initializer     ={type(self._initializer_base).__name__}\n'
-              f'    apply_layer_norm={self._normalize}')
+              f'    topology            ={self._topology}\n'
+              f'    activation_fn       ={self._activations[0].__name__}\n'
+              f'    initializer         ={type(self._initializer_base).__name__}\n'
+              f'    apply_input_norm    ={self._normalize}\n'
+              f'    input_norm_layerwise={self._normalize_per_layer}\n'
+              f'    input_norm_args     ={self._normalizer_kwargs}\n'
+              f'constraint-sat strategy:\n'
+              f'    parsed_action_bounds=\n        {bounds}\n'
+              f'    wrap_non_bool       ={self._wrap_non_bool}')
     
     def compile(self, compiled: JaxRDDLCompilerWithGrad,
-                _bounds: Dict, horizon: int) -> None:
+                _bounds: Bounds, 
+                horizon: int) -> None:
         rddl = compiled.rddl
         
         # calculate the correct action box bounds
@@ -737,7 +852,7 @@ class JaxDeepReactivePolicy(JaxPlan):
         if 1 < allowed_actions < bool_action_count:
             raise RDDLNotImplementedError(
                 f'Deep reactive policies currently do not support '
-                f'max-nondef-actions = {allowed_actions} > 1.')
+                f'max-nondef-actions {allowed_actions} > 1.')
         use_constraint_satisfaction = allowed_actions < bool_action_count
             
         noop = {var: (values[0] if isinstance(values, list) else values)
@@ -751,22 +866,75 @@ class JaxDeepReactivePolicy(JaxPlan):
                    
         ranges = rddl.variable_ranges
         normalize = self._normalize
+        normalize_per_layer = self._normalize_per_layer
+        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) 
                        for (var, shape) in shapes.items()}
         layer_names = {var: f'output_{var}'.replace('-', '_') for var in shapes}
         
-        # predict actions from the policy network for current state
-        def _jax_wrapped_policy_network_predict(state):
+        # inputs for the policy network
+        if rddl.observ_fluents:
+            observed_vars = rddl.observ_fluents
+        else:
+            observed_vars = rddl.state_fluents
+        input_names = {var: f'{var}'.replace('-', '_') for var in observed_vars}
+        
+        # catch if input norm is applied to size 1 tensor
+        if normalize:
+            non_bool_dims = 0
+            for (var, values) in observed_vars.items():
+                if ranges[var] != 'bool':
+                    value_size = np.atleast_1d(values).size
+                    if normalize_per_layer and value_size == 1:
+                        raise_warning(
+                            f'Cannot apply layer norm to state-fluent <{var}> '
+                            f'of size 1: setting normalize_per_layer = False.', 
+                            'red')
+                        normalize_per_layer = False
+                    non_bool_dims += value_size
+            if not normalize_per_layer and non_bool_dims == 1:
+                raise_warning(
+                    'Cannot apply layer norm to state-fluents of total size 1: '
+                    'setting normalize = False.', 'red')
+                normalize = False
+        
+        # convert subs dictionary into a state vector to feed to the MLP
+        def _jax_wrapped_policy_input(subs):
             
-            # apply layer norm
-            if normalize:
+            # concatenate all state variables into a single vector
+            # optionally apply layer norm to each input tensor
+            states_bool, states_non_bool = [], []
+            non_bool_dims = 0
+            for (var, value) in subs.items():
+                if var in observed_vars:
+                    state = jnp.ravel(value)
+                    if ranges[var] == 'bool':
+                        states_bool.append(state)
+                    else:
+                        if normalize and normalize_per_layer:
+                            normalizer = hk.LayerNorm(
+                                axis=-1, param_axis=-1, 
+                                name=f'input_norm_{input_names[var]}', 
+                                **self._normalizer_kwargs)
+                            state = normalizer(state)
+                        states_non_bool.append(state)
+                        non_bool_dims += state.size
+            state = jnp.concatenate(states_non_bool + states_bool)
+            
+            # 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,
-                    create_offset=True, create_scale=True,
-                    name='input_norm')
-                state = normalizer(state)
+                    axis=-1, param_axis=-1, name='input_norm', 
+                    **self._normalizer_kwargs)
+                normalized = normalizer(state[:non_bool_dims])
+                state = state.at[:non_bool_dims].set(normalized)
+            return state
+            
+        # predict actions from the policy network for current state
+        def _jax_wrapped_policy_network_predict(subs):
+            state = _jax_wrapped_policy_input(subs)
             
             # feed state vector through hidden layers
             hidden = state
@@ -789,16 +957,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,21 +997,14 @@ class JaxDeepReactivePolicy(JaxPlan):
                     actions[name] = action
                     start += size
             return actions
-                
-        # 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}
-            flat_subs = jax.tree_map(jnp.ravel, subs)
-            states = list(flat_subs.values())
-            state = jnp.concatenate(states)
-            return state
         
         # train action prediction
         def _jax_wrapped_drp_predict_train(key, params, hyperparams, step, subs):
-            state = _jax_wrapped_subs_to_state(subs)
-            actions = predict_fn.apply(params, state)
+            actions = predict_fn.apply(params, subs)
+            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 +1050,13 @@ 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}
-            state = _jax_wrapped_subs_to_state(subs)
-            params = predict_fn.init(key, state)
+                    if var in observed_vars}
+            params = predict_fn.init(key, subs)
             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 +1067,170 @@ 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, 
+                 show_violin: bool=True, show_action: bool=True) -> None:
+        '''Creates a new planner visualizer.
+        
+        :param rddl: the planning model to optimize
+        :param horizon: the lookahead or planning horizon
+        :param show_violin: whether to show the distribution of batch losses
+        :param show_action: whether to show heatmaps of the action fluents
+        '''
+        num_plots = 1
+        if show_violin:
+            num_plots += 1
+        if show_action:
+            num_plots += len(rddl.action_fluents)
+        self._fig, axes = plt.subplots(num_plots)
+        if num_plots == 1:
+            axes = [axes]
+        
+        # prepare the loss plot
+        self._loss_ax = axes[0]
+        self._loss_ax.autoscale(enable=True)
+        self._loss_ax.set_xlabel('training time')
+        self._loss_ax.set_ylabel('loss value')
+        self._loss_plot = self._loss_ax.plot(
+            [], [], linestyle=':', marker='o', markersize=2)[0]
+        self._loss_back = self._fig.canvas.copy_from_bbox(self._loss_ax.bbox)
+        
+        # prepare the violin plot
+        if show_violin:
+            self._hist_ax = axes[1]
+        else:
+            self._hist_ax = None
+            
+        # prepare the action plots
+        if show_action:
+            self._action_ax = {name: axes[idx + (2 if show_violin else 1)]
+                               for (idx, name) in enumerate(rddl.action_fluents)}
+            self._action_plots = {}
+            for name in rddl.action_fluents:
+                ax = self._action_ax[name]
+                if rddl.variable_ranges[name] == 'bool':
+                    vmin, vmax = 0.0, 1.0
+                else:
+                    vmin, vmax = None, None  
+                action_dim = 1
+                for dim in rddl.object_counts(rddl.variable_params[name]):
+                    action_dim *= dim     
+                action_plot = ax.pcolormesh(
+                    np.zeros((action_dim, horizon)), 
+                    cmap='seismic', vmin=vmin, vmax=vmax)
+                ax.set_aspect('auto')        
+                ax.set_xlabel('decision epoch')
+                ax.set_ylabel(name)
+                plt.colorbar(action_plot, ax=ax)
+                self._action_plots[name] = action_plot
+            self._action_back = {name: self._fig.canvas.copy_from_bbox(ax.bbox)
+                                 for (name, ax) in self._action_ax.items()}
+        else:
+            self._action_ax = None
+            self._action_plots = None
+            self._action_back = None
+            
+        plt.tight_layout()
+        plt.show(block=False)
+        
+    def redraw(self, xticks, losses, actions, returns) -> None:
+        
+        # draw the loss curve
+        self._fig.canvas.restore_region(self._loss_back)
+        self._loss_plot.set_xdata(xticks)
+        self._loss_plot.set_ydata(losses)        
+        self._loss_ax.set_xlim([0, len(xticks)])
+        self._loss_ax.set_ylim([np.min(losses), np.max(losses)])
+        self._loss_ax.draw_artist(self._loss_plot)
+        self._fig.canvas.blit(self._loss_ax.bbox)
+        
+        # draw the violin plot
+        if self._hist_ax is not None:
+            self._hist_ax.clear()
+            self._hist_ax.set_xlabel('loss value')
+            self._hist_ax.set_ylabel('density')
+            self._hist_ax.violinplot(returns, vert=False, showmeans=True)
+        
+        # draw the actions
+        if self._action_ax is not None:
+            for (name, values) in actions.items():
+                values = np.mean(values, axis=0, dtype=float)
+                values = np.reshape(values, newshape=(values.shape[0], -1)).T
+                self._fig.canvas.restore_region(self._action_back[name])
+                self._action_plots[name].set_array(values)
+                self._action_ax[name].draw_artist(self._action_plots[name])
+                self._fig.canvas.blit(self._action_ax[name].bbox)
+                self._action_plots[name].set_clim([np.min(values), np.max(values)])
+                
+        self._fig.canvas.draw()
+        self._fig.canvas.flush_events()
+        
+    def close(self) -> None:
+        plt.close(self._fig)
+        del self._loss_ax, self._hist_ax, self._action_ax, \
+            self._loss_plot, self._action_plots, self._fig, \
+            self._loss_back, self._action_back
+
+
+class JaxPlannerStatus(Enum):
+    '''Represents the status of a policy update from the JAX planner, 
+    including whether the update resulted in nan gradient, 
+    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
+    INVALID_GRADIENT = 3
+    TIME_BUDGET_REACHED = 4
+    ITER_BUDGET_REACHED = 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 +1255,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 +1275,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,33 +1302,84 @@ 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'
-              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'    cpfs_no_gradient={self.cpfs_without_grad}\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'optax {optax.__version__}, haiku {hk.__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'    user_action_bounds={self._action_bounds}\n'
+              f'    fuzzy logic type  ={type(self.logic).__name__}\n'
+              f'    nonfluents exact  ={self.compile_non_fluent_exact}\n'
+              f'    cpfs_no_gradient  ={self.cpfs_without_grad}\n'
               f'optimizer hyper-parameters:\n'
-              f'    use_64_bit      ={self.use64bit}\n'
-              f'    optimizer       ={self._optimizer_name.__name__}\n'
-              f'    optimizer args  ={self._optimizer_kwargs}\n'
-              f'    clip_gradient   ={self.clip_grad}\n'
-              f'    batch_size_train={self.batch_size_train}\n'
-              f'    batch_size_test ={self.batch_size_test}')
+              f'    use_64_bit        ={self.use64bit}\n'
+              f'    optimizer         ={self._optimizer_name.__name__}\n'
+              f'    optimizer args    ={self._optimizer_kwargs}\n'
+              f'    clip_gradient     ={self.clip_grad}\n'
+              f'    batch_size_train  ={self.batch_size_train}\n'
+              f'    batch_size_test   ={self.batch_size_test}')
         self.plan.summarize_hyperparameters()
         self.logic.summarize_hyperparameters()
         
+    # ===========================================================================
+    # COMPILATION SUBROUTINES
+    # ===========================================================================
+
     def _jax_compile_rddl(self):
         rddl = self.rddl
         
@@ -1017,13 +1387,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):
         
@@ -1039,6 +1414,7 @@ class JaxBackpropPlanner:
             policy=self.plan.train_policy,
             n_steps=self.horizon,
             n_batch=self.batch_size_train)
+        self.train_rollouts = train_rollouts
         
         test_rollouts = self.test_compiled.compile_rollouts(
             policy=self.plan.test_policy,
@@ -1051,11 +1427,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 +1443,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 +1459,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 +1472,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 +1483,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,15 +1522,106 @@ 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.
+    def as_optimization_problem(
+            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]:
+        '''Returns a function that computes the loss and a function that 
+        computes gradient of the return as a 1D vector given a 1D representation 
+        of policy parameters. These functions are designed to be compatible with 
+        off-the-shelf optimizers such as scipy. 
+        
+        Also returns the initial parameter vector to seed an optimizer, 
+        as well as a mapping that recovers the parameter pytree from the vector.
+        The PRNG key is updated internally starting from the optional given key.
+        
+        Constraints on actions, if they are required, cannot be constructed 
+        automatically in the general case. The user should build constraints
+        for each problem in the format required by the downstream optimizer.
+        
+        :param key: JAX PRNG key (derived from clock if not provided)
+        :param policy_hyperparameters: hyper-parameters for the policy/plan, 
+        such as weights for sigmoid wrapping boolean actions (defaults to 1
+        for all action-fluents if not provided)
+        :param loss_function_updates_key: if True, the loss function 
+        updates the PRNG key internally independently of the grad function
+        :param grad_function_updates_key: if True, the gradient function
+        updates the PRNG key internally independently of the loss function.
+        '''
         
-        :param key: JAX PRNG key   
+        # if PRNG key is not provided
+        if key is None:
+            key = random.PRNGKey(round(time.time() * 1000))
+            
+        # initialize the initial fluents, model parameters, policy hyper-params
+        subs = self.test_compiled.init_values
+        train_subs, _ = self._batched_init_subs(subs)
+        model_params = self.compiled.model_params
+        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}
+                
+        # initialize the policy parameters
+        params_guess, *_ = self.initialize(key, policy_hyperparams, train_subs)
+        guess_1d, unravel_fn = jax.flatten_util.ravel_pytree(params_guess)  
+        guess_1d = np.asarray(guess_1d)      
+        
+        # computes the training loss function and its 1D gradient
+        loss_fn = self._jax_loss(self.train_rollouts)
+        
+        @jax.jit
+        def _loss_with_key(key, params_1d):
+            policy_params = unravel_fn(params_1d)
+            loss_val, _ = loss_fn(key, policy_params, policy_hyperparams, 
+                                  train_subs, model_params)
+            return loss_val
+        
+        @jax.jit
+        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, 
+                                  train_subs, model_params)
+            grad_1d = jax.flatten_util.ravel_pytree(grad_val)[0]
+            return grad_1d
+        
+        def _loss_function(params_1d):
+            nonlocal key
+            if loss_function_updates_key:
+                key, subkey = random.split(key)
+            else:
+                subkey = key
+            loss_val = _loss_with_key(subkey, params_1d)
+            loss_val = float(loss_val)
+            return loss_val
+        
+        def _grad_function(params_1d):
+            nonlocal key
+            if grad_function_updates_key:
+                key, subkey = random.split(key)
+            else:
+                subkey = key
+            grad = _grad_with_key(subkey, params_1d)
+            grad = np.asarray(grad)
+            return grad
+        
+        return _loss_function, _grad_function, guess_1d, jax.jit(unravel_fn)
+        
+    # ===========================================================================
+    # 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 (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 plot_kwargs: additional arguments to pass to the plotter
         :param model_params: optional model-parameters to override default
         :param policy_hyperparams: hyper-parameters for the policy/plan, such as
         weights for sigmoid wrapping boolean actions
@@ -1161,64 +1629,110 @@ class JaxBackpropPlanner:
         their values: if None initializes all variables from the RDDL instance
         :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 print_summary: whether to print planner header, parameter 
+        summary, and diagnosis
+        :param print_progress: whether to print the progress bar during training
+        :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,
-                           verbose: int=2,
-                           tqdm_position: int=None) -> Generator[Dict[str, object], None, None]:
-        '''Returns a generator for computing an optimal straight-line plan. 
+                           plot_step: Optional[int]=None,
+                           plot_kwargs: Optional[Dict[str, Any]]=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,
+                           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 plot_kwargs: additional arguments to pass to the plotter
         :param model_params: optional model-parameters to override default
         :param policy_hyperparams: hyper-parameters for the policy/plan, such as
         weights for sigmoid wrapping boolean actions
         :param subs: dictionary mapping initial state and non-fluents to 
         their values: if None initializes all variables from the RDDL instance
         :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)
+        specified in this instance        
+        :param print_summary: whether to print planner header, parameter 
+        summary, and diagnosis
+        :param print_progress: whether to print the progress bar during training
+        :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'
-                  '==============================================')
+        if print_summary:
+            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'    plot_kwargs        ={plot_kwargs}\n'
+                  f'    print_summary      ={print_summary}\n'
+                  f'    print_progress     ={print_progress}\n')
+            if self.compiled.relaxations:
+                print('Some RDDL operations are non-differentiable, '
+                      'replacing them with differentiable relaxations:')
+                print(self.compiled.summarize_model_relaxations())
             
         # compute a batched version of the initial values
         if subs is None:
@@ -1237,7 +1751,7 @@ class JaxBackpropPlanner:
                               'from the RDDL files.')
         train_subs, test_subs = self._batched_init_subs(subs)
         
-        # initialize, model parameters
+        # initialize model parameters
         if model_params is None:
             model_params = self.compiled.model_params
         model_params_test = self.test_compiled.model_params
@@ -1245,63 +1759,103 @@ 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:
+            if plot_kwargs is None:
+                plot_kwargs = {}
+            plot = JaxPlannerPlot(self.rddl, self.horizon, **plot_kwargs)
+        xticks, loss_values = [], []
         
         # training loop
         iters = range(epochs)
-        if verbose >= 2:
+        if print_progress:
             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)
+            
+            # 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
+             
+            # constraint satisfaction problem
             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)
+            # 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
+              
+            # evaluate test losses and record best plan so far
             test_loss, log = self.test_loss(
-                subkey3, policy_params, policy_hyperparams,
+                subkey, policy_params, policy_hyperparams,
                 test_subs, model_params_test)
-            
-            # record the best plan so far
+            test_loss = rolling_test_loss.update(test_loss)
             if test_loss < best_loss:
                 best_params, best_loss, best_grad = \
                     policy_params, test_loss, train_log['grad']
                 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}
+                returns = -np.sum(np.asarray(log['reward']), axis=1)
+                plot.redraw(xticks, loss_values, action_values, returns)
             
             # if the progress bar is used
             elapsed = time.time() - start_time - elapsed_outside_loop
-            if verbose >= 2:
+            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.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
             
             # return a callback
             start_time_outside = time.time()
             yield {
+                'status': status,
                 'iteration': it,
                 'train_return':-train_loss,
                 'test_return':-test_loss,
@@ -1318,16 +1872,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
-        
-        if verbose >= 2:
+                        
+        # release resources
+        if print_progress:
             iters.close()
+        if plot is not None:
+            plot.close()
         
         # validate the test return
         if log:
@@ -1337,24 +1890,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)
+        if print_summary:
+            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 +1925,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 +1949,213 @@ 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'make sure vectorized = True in the RDDLEnv.')
+            
+            # must be numeric array
+            # exception is for POMDPs at 1st epoch when observ-fluents are None
+            dtype = np.atleast_1d(values).dtype
+            if not jnp.issubdtype(dtype, jnp.number) \
+            and not jnp.issubdtype(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 {values} assigned to p-variable <{var}> are '
+                        f'non-numeric of type {dtype}.')
+            
         # 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,
+                 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__(
-            *args,
-            optimizer=optimizer,
-            optimizer_kwargs=optimizer_kwargs,
-            **kwargs)
-        
-    def summarize_hyperparameters(self):
-        super(JaxArmijoLineSearchPlanner, self).summarize_hyperparameters()
+        self.step_max = step_max
+        self.step_min = step_min
+        if 'clip_grad' in kwargs:
+            raise_warning('clip_grad parameter conflicts with '
+                          'line search planner and will be ignored.', 'red')
+            del kwargs['clip_grad']
+        super(JaxLineSearchPlanner, self).__init__(*args, **kwargs)
+        
+    def summarize_hyperparameters(self) -> None:
+        super(JaxLineSearchPlanner, self).summarize_hyperparameters()
         print(f'linesearch hyper-parameters:\n'
-              f'    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) \
+            or not trials:
+                trials += 1
+                step *= decay
+                f_step, new_params, new_state = _jax_wrapped_line_search_trial(
+                    step, grad, key, policy_params, hyperparams, subs, 
+                    model_params, opt_state)
+                if f_step < best_f:
+                    best_f, best_step, best_params, best_state = \
+                        f_step, step, new_params, new_state
             
-            # 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) - 0.5 * beta * 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 +2165,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,60 +2175,72 @@ 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)
+            callback = self.planner.optimize(key=self.key, **self.train_kwargs)
+            self.params = callback['best_params']
 
 
 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