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

pyRDDLGym_jax/core/planner.py

@@ -47,7 +47,6 @@ Bounds = Dict[str, Tuple[np.ndarray, np.ndarray]]
 Kwargs = Dict[str, Any]
 Pytree = Any
 
-
 # ***********************************************************************
 # CONFIG FILE MANAGEMENT
 # 
@@ -57,6 +56,7 @@ Pytree = Any
 #
 # ***********************************************************************
 
+
 def _parse_config_file(path: str):
     if not os.path.isfile(path):
         raise FileNotFoundError(f'File {path} does not exist.')
@@ -103,10 +103,13 @@ def _load_config(config, args):
     compare_kwargs = model_args.get('comparison_kwargs', {})
     sampling_name = model_args.get('sampling', 'GumbelSoftmax')
     sampling_kwargs = model_args.get('sampling_kwargs', {})
+    rounding_name = model_args.get('rounding', 'SoftRounding')
+    rounding_kwargs = model_args.get('rounding_kwargs', {})
     logic_kwargs['tnorm'] = getattr(logic, tnorm_name)(**tnorm_kwargs)
     logic_kwargs['complement'] = getattr(logic, comp_name)(**comp_kwargs)
     logic_kwargs['comparison'] = getattr(logic, compare_name)(**compare_kwargs)
     logic_kwargs['sampling'] = getattr(logic, sampling_name)(**sampling_kwargs)
+    logic_kwargs['rounding'] = getattr(logic, rounding_name)(**rounding_kwargs)
     
     # read the policy settings
     plan_method = planner_args.pop('method')
@@ -157,11 +160,18 @@ def _load_config(config, args):
         else:
             planner_args['optimizer'] = optimizer
         
-    # read the optimize call settings
+    # optimize call RNG key
     planner_key = train_args.get('key', None)
     if planner_key is not None:
         train_args['key'] = random.PRNGKey(planner_key)
     
+    # optimize call stopping rule
+    stopping_rule = train_args.get('stopping_rule', None)
+    if stopping_rule is not None:
+        stopping_rule_kwargs = train_args.pop('stopping_rule_kwargs', {})
+        train_args['stopping_rule'] = getattr(
+            sys.modules[__name__], stopping_rule)(**stopping_rule_kwargs)
+    
     return planner_args, plan_kwargs, train_args
 
 
@@ -175,7 +185,6 @@ def load_config_from_string(value: str) -> Tuple[Kwargs, ...]:
     '''Loads config file contents specified explicitly as a string value.'''
     config, args = _parse_config_string(value)
     return _load_config(config, args)
-
     
 # ***********************************************************************
 # MODEL RELAXATIONS
@@ -299,7 +308,6 @@ class JaxRDDLCompilerWithGrad(JaxRDDLCompiler):
         arg, = expr.args
         arg = self._jax(arg, info)
         return arg
-    
 
 # ***********************************************************************
 # ALL VERSIONS OF JAX PLANS
@@ -309,6 +317,7 @@ class JaxRDDLCompilerWithGrad(JaxRDDLCompiler):
 #
 # ***********************************************************************
 
+
 class JaxPlan:
     '''Base class for all JAX policy representations.'''
     
@@ -363,7 +372,7 @@ class JaxPlan:
         self._projection = value
     
     def _calculate_action_info(self, compiled: JaxRDDLCompilerWithGrad,
-                               user_bounds: Bounds, 
+                               user_bounds: Bounds,
                                horizon: int):
         shapes, bounds, bounds_safe, cond_lists = {}, {}, {}, {}
         for (name, prange) in compiled.rddl.variable_ranges.items():
@@ -463,7 +472,7 @@ class JaxStraightLinePlan(JaxPlan):
               f'    max_projection_iters ={self._max_constraint_iter}')
     
     def compile(self, compiled: JaxRDDLCompilerWithGrad,
-                _bounds: Bounds, 
+                _bounds: Bounds,
                 horizon: int) -> None:
         rddl = compiled.rddl
         
@@ -504,7 +513,7 @@ class JaxStraightLinePlan(JaxPlan):
         def _jax_bool_action_to_param(var, action, hyperparams):
             if wrap_sigmoid:
                 weight = hyperparams[var]
-                return (-1.0 / weight) * jnp.log(1.0 / action - 1.0)
+                return jax.scipy.special.logit(action) / weight
             else:
                 return action
             
@@ -513,14 +522,13 @@ class JaxStraightLinePlan(JaxPlan):
         def _jax_non_bool_param_to_action(var, param, hyperparams):
             if wrap_non_bool:
                 lower, upper = bounds_safe[var]
-                action = jnp.select(
-                    condlist=cond_lists[var],
-                    choicelist=[
-                        lower + (upper - lower) * jax.nn.sigmoid(param),
-                        lower + (jax.nn.elu(param) + 1.0),
-                        upper - (jax.nn.elu(-param) + 1.0),
-                        param
-                    ]
+                mb, ml, mu, mn = [mask.astype(compiled.REAL) 
+                                  for mask in cond_lists[var]]       
+                action = (
+                    mb * (lower + (upper - lower) * jax.nn.sigmoid(param)) + 
+                    ml * (lower + (jax.nn.elu(param) + 1.0)) + 
+                    mu * (upper - (jax.nn.elu(-param) + 1.0)) + 
+                    mn * param
                 )
             else:
                 action = param
@@ -780,7 +788,7 @@ class JaxDeepReactivePolicy(JaxPlan):
     def __init__(self, topology: Optional[Sequence[int]]=None,
                  activation: Activation=jnp.tanh,
                  initializer: hk.initializers.Initializer=hk.initializers.VarianceScaling(scale=2.0),
-                 normalize: bool=False, 
+                 normalize: bool=False,
                  normalize_per_layer: bool=False,
                  normalizer_kwargs: Optional[Kwargs]=None,
                  wrap_non_bool: bool=False) -> None:
@@ -828,7 +836,7 @@ class JaxDeepReactivePolicy(JaxPlan):
               f'    wrap_non_bool       ={self._wrap_non_bool}')
     
     def compile(self, compiled: JaxRDDLCompilerWithGrad,
-                _bounds: Bounds, 
+                _bounds: Bounds,
                 horizon: int) -> None:
         rddl = compiled.rddl
         
@@ -881,7 +889,7 @@ class JaxDeepReactivePolicy(JaxPlan):
                     if normalize_per_layer and value_size == 1:
                         raise_warning(
                             f'Cannot apply layer norm to state-fluent <{var}> '
-                            f'of size 1: setting normalize_per_layer = False.', 
+                            f'of size 1: setting normalize_per_layer = False.',
                             'red')
                         normalize_per_layer = False
                     non_bool_dims += value_size
@@ -906,8 +914,8 @@ class JaxDeepReactivePolicy(JaxPlan):
                     else:
                         if normalize and normalize_per_layer:
                             normalizer = hk.LayerNorm(
-                                axis=-1, param_axis=-1, 
-                                name=f'input_norm_{input_names[var]}', 
+                                axis=-1, param_axis=-1,
+                                name=f'input_norm_{input_names[var]}',
                                 **self._normalizer_kwargs)
                             state = normalizer(state)
                         states_non_bool.append(state)
@@ -917,7 +925,7 @@ class JaxDeepReactivePolicy(JaxPlan):
             # optionally perform layer normalization on the non-bool inputs
             if normalize and not normalize_per_layer and non_bool_dims:
                 normalizer = hk.LayerNorm(
-                    axis=-1, param_axis=-1, name='input_norm', 
+                    axis=-1, param_axis=-1, name='input_norm',
                     **self._normalizer_kwargs)
                 normalized = normalizer(state[:non_bool_dims])
                 state = state.at[:non_bool_dims].set(normalized)
@@ -950,14 +958,13 @@ class JaxDeepReactivePolicy(JaxPlan):
                 else:
                     if wrap_non_bool:
                         lower, upper = bounds_safe[var]
-                        action = jnp.select(
-                            condlist=cond_lists[var],
-                            choicelist=[
-                                lower + (upper - lower) * jax.nn.sigmoid(output),
-                                lower + (jax.nn.elu(output) + 1.0),
-                                upper - (jax.nn.elu(-output) + 1.0),
-                                output
-                            ]
+                        mb, ml, mu, mn = [mask.astype(compiled.REAL) 
+                                          for mask in cond_lists[var]]       
+                        action = (
+                            mb * (lower + (upper - lower) * jax.nn.sigmoid(output)) + 
+                            ml * (lower + (jax.nn.elu(output) + 1.0)) + 
+                            mu * (upper - (jax.nn.elu(-output) + 1.0)) + 
+                            mn * output
                         )
                     else:
                         action = output
@@ -1049,7 +1056,6 @@ class JaxDeepReactivePolicy(JaxPlan):
         
     def guess_next_epoch(self, params: Pytree) -> Pytree:
         return params
-
     
 # ***********************************************************************
 # ALL VERSIONS OF JAX PLANNER
@@ -1059,6 +1065,7 @@ class JaxDeepReactivePolicy(JaxPlan):
 #
 # ***********************************************************************
 
+
 class RollingMean:
     '''Maintains an estimate of the rolling mean of a stream of real-valued 
     observations.'''
@@ -1080,7 +1087,7 @@ class RollingMean:
 class JaxPlannerPlot:
     '''Supports plotting and visualization of a JAX policy in real time.'''
     
-    def __init__(self, rddl: RDDLPlanningModel, horizon: int, 
+    def __init__(self, rddl: RDDLPlanningModel, horizon: int,
                  show_violin: bool=True, show_action: bool=True) -> None:
         '''Creates a new planner visualizer.
         
@@ -1128,7 +1135,7 @@ class JaxPlannerPlot:
                 for dim in rddl.object_counts(rddl.variable_params[name]):
                     action_dim *= dim     
                 action_plot = ax.pcolormesh(
-                    np.zeros((action_dim, horizon)), 
+                    np.zeros((action_dim, horizon)),
                     cmap='seismic', vmin=vmin, vmax=vmax)
                 ax.set_aspect('auto')        
                 ax.set_xlabel('decision epoch')
@@ -1201,6 +1208,39 @@ class JaxPlannerStatus(Enum):
         return self.value >= 3
 
 
+class JaxPlannerStoppingRule:
+    '''The base class of all planner stopping rules.'''
+    
+    def reset(self) -> None:
+        raise NotImplementedError
+        
+    def monitor(self, callback: Dict[str, Any]) -> bool:
+        raise NotImplementedError
+    
+
+class NoImprovementStoppingRule(JaxPlannerStoppingRule):
+    '''Stopping rule based on no improvement for a fixed number of iterations.'''
+    
+    def __init__(self, patience: int) -> None:
+        self.patience = patience
+    
+    def reset(self) -> None:
+        self.callback = None
+        self.iters_since_last_update = 0
+        
+    def monitor(self, callback: Dict[str, Any]) -> bool:
+        if self.callback is None \
+        or callback['best_return'] > self.callback['best_return']:
+            self.callback = callback
+            self.iters_since_last_update = 0
+        else:
+            self.iters_since_last_update += 1
+        return self.iters_since_last_update >= self.patience
+    
+    def __str__(self) -> str:
+        return f'No improvement for {self.patience} iterations'
+        
+
 class JaxBackpropPlanner:
     '''A class for optimizing an action sequence in the given RDDL MDP using 
     gradient descent.'''
@@ -1215,6 +1255,8 @@ class JaxBackpropPlanner:
                  optimizer: Callable[..., optax.GradientTransformation]=optax.rmsprop,
                  optimizer_kwargs: Optional[Kwargs]=None,
                  clip_grad: Optional[float]=None,
+                 noise_grad_eta: float=0.0,
+                 noise_grad_gamma: float=1.0,
                  logic: FuzzyLogic=FuzzyLogic(),
                  use_symlog_reward: bool=False,
                  utility: Union[Callable[[jnp.ndarray], float], str]='mean',
@@ -1241,6 +1283,8 @@ class JaxBackpropPlanner:
         :param optimizer_kwargs: a dictionary of parameters to pass to the SGD
         factory (e.g. which parameters are controllable externally)
         :param clip_grad: maximum magnitude of gradient updates
+        :param noise_grad_eta: scale of the gradient noise variance
+        :param noise_grad_gamma: decay rate of the gradient noise variance
         :param logic: a subclass of FuzzyLogic for mapping exact mathematical
         operations to their differentiable counterparts 
         :param use_symlog_reward: whether to use the symlog transform on the 
@@ -1275,6 +1319,8 @@ class JaxBackpropPlanner:
             optimizer_kwargs = {'learning_rate': 0.1}
         self._optimizer_kwargs = optimizer_kwargs
         self.clip_grad = clip_grad
+        self.noise_grad_eta = noise_grad_eta
+        self.noise_grad_gamma = noise_grad_gamma
         
         # set optimizer
         try:
@@ -1340,14 +1386,14 @@ class JaxBackpropPlanner:
         except Exception as _:
             devices_short = 'N/A'
         LOGO = \
+r"""
+   __   ______   __  __   ______  __       ______   __   __    
+  /\ \ /\  __ \ /\_\_\_\ /\  == \/\ \     /\  __ \ /\ "-.\ \   
+ _\_\ \\ \  __ \\/_/\_\/_\ \  _-/\ \ \____\ \  __ \\ \ \-.  \  
+/\_____\\ \_\ \_\ /\_\/\_\\ \_\   \ \_____\\ \_\ \_\\ \_\\"\_\ 
+\/_____/ \/_/\/_/ \/_/\/_/ \/_/    \/_____/ \/_/\/_/ \/_/ \/_/ 
 """
-   __    ______    __  __    ______  __        ______    __   __    
-  /\ \  /\  __ \  /\_\_\_\  /\  == \/\ \      /\  __ \  /\ "-.\ \   
- _\_\ \ \ \  __ \ \/_/\_\/_ \ \  _-/\ \ \____ \ \  __ \ \ \ \-.  \  
-/\_____\ \ \_\ \_\  /\_\/\_\ \ \_\   \ \_____\ \ \_\ \_\ \ \_\\"\_\ 
-\/_____/  \/_/\/_/  \/_/\/_/  \/_/    \/_____/  \/_/\/_/  \/_/ \/_/ 
-"""
-                                                       
+                   
         print('\n'
               f'{LOGO}\n'
               f'Version {__version__}\n' 
@@ -1372,6 +1418,8 @@ class JaxBackpropPlanner:
               f'    optimizer         ={self._optimizer_name.__name__}\n'
               f'    optimizer args    ={self._optimizer_kwargs}\n'
               f'    clip_gradient     ={self.clip_grad}\n'
+              f'    noise_grad_eta    ={self.noise_grad_eta}\n'
+              f'    noise_grad_gamma  ={self.noise_grad_gamma}\n'
               f'    batch_size_train  ={self.batch_size_train}\n'
               f'    batch_size_test   ={self.batch_size_test}')
         self.plan.summarize_hyperparameters()
@@ -1396,7 +1444,7 @@ class JaxBackpropPlanner:
         
         # Jax compilation of the exact RDDL for testing
         self.test_compiled = JaxRDDLCompiler(
-            rddl=rddl, 
+            rddl=rddl,
             logger=self.logger,
             use64bit=self.use64bit)
         self.test_compiled.compile(log_jax_expr=True, heading='EXACT MODEL')
@@ -1473,7 +1521,7 @@ class JaxBackpropPlanner:
         def _jax_wrapped_init_policy(key, hyperparams, subs):
             policy_params = init(key, hyperparams, subs)
             opt_state = optimizer.init(policy_params)
-            return policy_params, opt_state, None
+            return policy_params, opt_state, {}
         
         return _jax_wrapped_init_policy
         
@@ -1481,6 +1529,19 @@ class JaxBackpropPlanner:
         optimizer = self.optimizer
         projection = self.plan.projection
         
+        # add Gaussian gradient noise per Neelakantan et al., 2016.
+        def _jax_wrapped_gaussian_param_noise(key, grads, sigma):
+            treedef = jax.tree_util.tree_structure(grads)
+            keys_flat = random.split(key, num=treedef.num_leaves)            
+            keys_tree = jax.tree_util.tree_unflatten(treedef, keys_flat)
+            new_grads = jax.tree_map(
+                lambda g, k: g + sigma * random.normal(
+                    key=k, shape=g.shape, dtype=g.dtype),
+                grads,
+                keys_tree
+            )
+            return new_grads
+        
         # calculate the plan gradient w.r.t. return loss and update optimizer
         # also perform a projection step to satisfy constraints on actions
         def _jax_wrapped_plan_update(key, policy_params, hyperparams,
@@ -1488,12 +1549,14 @@ class JaxBackpropPlanner:
             grad_fn = jax.value_and_grad(loss, argnums=1, has_aux=True)
             (loss_val, log), grad = grad_fn(
                 key, policy_params, hyperparams, subs, model_params)  
+            sigma = opt_aux.get('noise_sigma', 0.0)
+            grad = _jax_wrapped_gaussian_param_noise(key, grad, sigma)
             updates, opt_state = optimizer.update(grad, opt_state) 
             policy_params = optax.apply_updates(policy_params, updates)
             policy_params, converged = projection(policy_params, hyperparams)
             log['grad'] = grad
             log['updates'] = updates
-            return policy_params, converged, opt_state, None, loss_val, log
+            return policy_params, converged, opt_state, opt_aux, loss_val, log
         
         return jax.jit(_jax_wrapped_plan_update)
             
@@ -1524,7 +1587,7 @@ class JaxBackpropPlanner:
         return init_train, init_test
     
     def as_optimization_problem(
-            self, key: Optional[random.PRNGKey]=None, 
+            self, key: Optional[random.PRNGKey]=None,
             policy_hyperparams: Optional[Pytree]=None,
             loss_function_updates_key: bool=True,
             grad_function_updates_key: bool=False) -> Tuple[Callable, Callable, np.ndarray, Callable]:
@@ -1576,7 +1639,7 @@ class JaxBackpropPlanner:
         @jax.jit
         def _loss_with_key(key, params_1d):
             policy_params = unravel_fn(params_1d)
-            loss_val, _ = loss_fn(key, policy_params, policy_hyperparams, 
+            loss_val, _ = loss_fn(key, policy_params, policy_hyperparams,
                                   train_subs, model_params)
             return loss_val
         
@@ -1584,7 +1647,7 @@ class JaxBackpropPlanner:
         def _grad_with_key(key, params_1d):
             policy_params = unravel_fn(params_1d)
             grad_fn = jax.grad(loss_fn, argnums=1, has_aux=True)
-            grad_val, _ = grad_fn(key, policy_params, policy_hyperparams, 
+            grad_val, _ = grad_fn(key, policy_params, policy_hyperparams,
                                   train_subs, model_params)
             grad_1d = jax.flatten_util.ravel_pytree(grad_val)[0]
             return grad_1d
@@ -1633,6 +1696,7 @@ class JaxBackpropPlanner:
         :param print_summary: whether to print planner header, parameter 
         summary, and diagnosis
         :param print_progress: whether to print the progress bar during training
+        :param stopping_rule: stopping criterion
         :param test_rolling_window: the test return is averaged on a rolling 
         window of the past test_rolling_window returns when updating the best
         parameters found so far
@@ -1658,13 +1722,14 @@ class JaxBackpropPlanner:
                            epochs: int=999999,
                            train_seconds: float=120.,
                            plot_step: Optional[int]=None,
-                           plot_kwargs: Optional[Dict[str, Any]]=None,
+                           plot_kwargs: Optional[Kwargs]=None,
                            model_params: Optional[Dict[str, Any]]=None,
                            policy_hyperparams: Optional[Dict[str, Any]]=None,
                            subs: Optional[Dict[str, Any]]=None,
                            guess: Optional[Pytree]=None,
                            print_summary: bool=True,
                            print_progress: bool=True,
+                           stopping_rule: Optional[JaxPlannerStoppingRule]=None,
                            test_rolling_window: int=10,
                            tqdm_position: Optional[int]=None) -> Generator[Dict[str, Any], None, None]:
         '''Returns a generator for computing an optimal policy or plan. 
@@ -1686,6 +1751,7 @@ class JaxBackpropPlanner:
         :param print_summary: whether to print planner header, parameter 
         summary, and diagnosis
         :param print_progress: whether to print the progress bar during training
+        :param stopping_rule: stopping criterion
         :param test_rolling_window: the test return is averaged on a rolling 
         window of the past test_rolling_window returns when updating the best
         parameters found so far
@@ -1737,10 +1803,11 @@ class JaxBackpropPlanner:
                   f'    plot_frequency     ={plot_step}\n'
                   f'    plot_kwargs        ={plot_kwargs}\n'
                   f'    print_summary      ={print_summary}\n'
-                  f'    print_progress     ={print_progress}\n')
+                  f'    print_progress     ={print_progress}\n'
+                  f'    stopping_rule      ={stopping_rule}\n')
             if self.compiled.relaxations:
                 print('Some RDDL operations are non-differentiable, '
-                      'replacing them with differentiable relaxations:')
+                      'they will be approximated as follows:')
                 print(self.compiled.summarize_model_relaxations())
             
         # compute a batched version of the initial values
@@ -1773,7 +1840,7 @@ class JaxBackpropPlanner:
         else:
             policy_params = guess
             opt_state = self.optimizer.init(policy_params)
-            opt_aux = None
+            opt_aux = {}
         
         # initialize running statistics
         best_params, best_loss, best_grad = policy_params, jnp.inf, jnp.inf
@@ -1783,6 +1850,10 @@ class JaxBackpropPlanner:
         status = JaxPlannerStatus.NORMAL
         is_all_zero_fn = lambda x: np.allclose(x, 0)
         
+        # initialize stopping criterion
+        if stopping_rule is not None:
+            stopping_rule.reset()
+            
         # initialize plot area
         if plot_step is None or plot_step <= 0 or plt is None:
             plot = None
@@ -1801,6 +1872,11 @@ class JaxBackpropPlanner:
         for it in iters:
             status = JaxPlannerStatus.NORMAL
             
+            # gradient noise schedule
+            noise_var = self.noise_grad_eta / (1. + it) ** self.noise_grad_gamma
+            noise_sigma = np.sqrt(noise_var)
+            opt_aux['noise_sigma'] = noise_sigma
+            
             # update the parameters of the plan
             key, subkey = random.split(key)
             policy_params, converged, opt_state, opt_aux, \
@@ -1865,8 +1941,7 @@ class JaxBackpropPlanner:
                 status = JaxPlannerStatus.ITER_BUDGET_REACHED
             
             # return a callback
-            start_time_outside = time.time()
-            yield {
+            callback = {
                 'status': status,
                 'iteration': it,
                 'train_return':-train_loss,
@@ -1877,16 +1952,23 @@ class JaxBackpropPlanner:
                 'last_iteration_improved': last_iter_improve,
                 'grad': train_log['grad'],
                 'best_grad': best_grad,
+                'noise_sigma': noise_sigma,
                 'updates': train_log['updates'],
                 'elapsed_time': elapsed,
                 'key': key,
                 **log
             }
+            start_time_outside = time.time()
+            yield callback
             elapsed_outside_loop += (time.time() - start_time_outside)
             
             # abortion check
             if status.is_failure():
                 break
+            
+            # stopping condition reached
+            if stopping_rule is not None and stopping_rule.monitor(callback):
+                break                
                         
         # release resources
         if print_progress:
@@ -1918,7 +2000,7 @@ class JaxBackpropPlanner:
                   f'    best_grad_norm={grad_norm}\n'
                   f'    diagnosis: {diagnosis}\n')
     
-    def _perform_diagnosis(self, last_iter_improve, 
+    def _perform_diagnosis(self, last_iter_improve,
                            train_return, test_return, best_return, grad_norm):
         max_grad_norm = max(jax.tree_util.tree_leaves(grad_norm))
         grad_is_zero = np.allclose(max_grad_norm, 0)
@@ -2097,7 +2179,7 @@ class JaxLineSearchPlanner(JaxBackpropPlanner):
                 trials += 1
                 step *= decay
                 f_step, new_params, new_state = _jax_wrapped_line_search_trial(
-                    step, grad, key, policy_params, hyperparams, subs, 
+                    step, grad, key, policy_params, hyperparams, subs,
                     model_params, opt_state)
                 if f_step < best_f:
                     best_f, best_step, best_params, best_state = \
@@ -2106,11 +2188,11 @@ class JaxLineSearchPlanner(JaxBackpropPlanner):
             log['updates'] = None
             log['line_search_iters'] = trials
             log['learning_rate'] = best_step
-            return best_params, True, best_state, best_step, best_f, log
+            opt_aux['best_step'] = best_step
+            return best_params, True, best_state, opt_aux, best_f, log
             
         return _jax_wrapped_plan_update
 
-
 # ***********************************************************************
 # ALL VERSIONS OF RISK FUNCTIONS
 # 
@@ -2141,7 +2223,6 @@ def cvar_utility(returns: jnp.ndarray, alpha: float) -> float:
     alpha_mask = jax.lax.stop_gradient(
         returns <= jnp.percentile(returns, q=100 * alpha))
     return jnp.sum(returns * alpha_mask) / jnp.sum(alpha_mask)
-    
 
 # ***********************************************************************
 # ALL VERSIONS OF CONTROLLERS
@@ -2151,12 +2232,13 @@ def cvar_utility(returns: jnp.ndarray, alpha: float) -> float:
 #
 # ***********************************************************************
 
+
 class JaxOfflineController(BaseAgent):
     '''A container class for a Jax policy trained offline.'''
     
     use_tensor_obs = True
     
-    def __init__(self, planner: JaxBackpropPlanner, 
+    def __init__(self, planner: JaxBackpropPlanner,
                  key: Optional[random.PRNGKey]=None,
                  eval_hyperparams: Optional[Dict[str, Any]]=None,
                  params: Optional[Pytree]=None,
@@ -2211,7 +2293,7 @@ class JaxOnlineController(BaseAgent):
     
     use_tensor_obs = True
     
-    def __init__(self, planner: JaxBackpropPlanner, 
+    def __init__(self, planner: JaxBackpropPlanner,
                  key: Optional[random.PRNGKey]=None,
                  eval_hyperparams: Optional[Dict[str, Any]]=None,
                  warm_start: bool=True,