[tune] Make PBT Quantile fraction configurable (#4912)

python/ray/tune/schedulers/pbt.py

@@ -19,10 +19,6 @@ from ray.tune.trial import Trial, Checkpoint
 
 logger = logging.getLogger(__name__)
 
-# Parameters are transferred from the top PBT_QUANTILE fraction of trials to
-# the bottom PBT_QUANTILE fraction.
-PBT_QUANTILE = 0.25
-
 
 class PBTTrialState(object):
     """Internal PBT state tracked per-trial."""
@@ -134,6 +130,10 @@ class PopulationBasedTraining(FIFOScheduler):
             A function specifies the distribution of a continuous parameter.
             You must specify at least one of `hyperparam_mutations` or
             `custom_explore_fn`.
+        quantile_fraction (float): Parameters are transferred from the top
+            `quantile_fraction` fraction of trials to the bottom
+            `quantile_fraction` fraction. Needs to be between 0 and 0.5.
+            Setting it to 0 essentially implies doing no exploitation at all.
         resample_probability (float): The probability of resampling from the
             original distribution when applying `hyperparam_mutations`. If not
             resampled, the value will be perturbed by a factor of 1.2 or 0.8
@@ -172,6 +172,7 @@ class PopulationBasedTraining(FIFOScheduler):
                  mode="max",
                  perturbation_interval=60.0,
                  hyperparam_mutations={},
+                 quantile_fraction=0.25,
                  resample_probability=0.25,
                  custom_explore_fn=None,
                  log_config=True):
@@ -180,6 +181,11 @@ class PopulationBasedTraining(FIFOScheduler):
                 "You must specify at least one of `hyperparam_mutations` or "
                 "`custom_explore_fn` to use PBT.")
 
+        if quantile_fraction > 0.5 or quantile_fraction < 0:
+            raise TuneError(
+                "You must set `quantile_fraction` to a value between 0 and"
+                "0.5. Current value: '{}'".format(quantile_fraction))
+
         assert mode in ["min", "max"], "`mode` must be 'min' or 'max'!"
 
         if reward_attr is not None:
@@ -199,6 +205,7 @@ class PopulationBasedTraining(FIFOScheduler):
         self._time_attr = time_attr
         self._perturbation_interval = perturbation_interval
         self._hyperparam_mutations = hyperparam_mutations
+        self._quantile_fraction = quantile_fraction
         self._resample_probability = resample_probability
         self._trial_state = {}
         self._custom_explore_fn = custom_explore_fn
@@ -247,6 +254,7 @@ class PopulationBasedTraining(FIFOScheduler):
 
         For each step, logs: [target trial tag, clone trial tag, target trial
         iteration, clone trial iteration, old config, new config].
+
         """
         trial_name, trial_to_clone_name = (trial_state.orig_tag,
                                            new_state.orig_tag)
@@ -277,7 +285,9 @@ class PopulationBasedTraining(FIFOScheduler):
     def _exploit(self, trial_executor, trial, trial_to_clone):
         """Transfers perturbed state from trial_to_clone -> trial.
 
-        If specified, also logs the updated hyperparam state."""
+        If specified, also logs the updated hyperparam state.
+
+        """
 
         trial_state = self._trial_state[trial]
         new_state = self._trial_state[trial_to_clone]
@@ -318,7 +328,9 @@ class PopulationBasedTraining(FIFOScheduler):
     def _quantiles(self):
         """Returns trials in the lower and upper `quantile` of the population.
 
-        If there is not enough data to compute this, returns empty lists."""
+        If there is not enough data to compute this, returns empty lists.
+
+        """
 
         trials = []
         for trial, state in self._trial_state.items():
@@ -329,14 +341,19 @@ class PopulationBasedTraining(FIFOScheduler):
         if len(trials) <= 1:
             return [], []
         else:
-            return (trials[:int(math.ceil(len(trials) * PBT_QUANTILE))],
-                    trials[int(math.floor(-len(trials) * PBT_QUANTILE)):])
+            num_trials_in_quantile = int(
+                math.ceil(len(trials) * self._quantile_fraction))
+            if num_trials_in_quantile > len(trials) / 2:
+                num_trials_in_quantile = int(math.floor(len(trials) / 2))
+            return (trials[:num_trials_in_quantile],
+                    trials[-num_trials_in_quantile:])
 
     def choose_trial_to_run(self, trial_runner):
         """Ensures all trials get fair share of time (as defined by time_attr).
 
         This enables the PBT scheduler to support a greater number of
         concurrent trials than can fit in the cluster at any given time.
+
         """
 
         candidates = []

python/ray/tune/tests/test_trial_scheduler.py

@@ -627,6 +627,7 @@ class PopulationBasedTestingSuite(unittest.TestCase):
             time_attr="training_iteration",
             perturbation_interval=10,
             resample_probability=resample_prob,
+            quantile_fraction=0.25,
             hyperparam_mutations={
                 "id_factor": [100],
                 "float_factor": lambda: 100.0,