Implement early stopping and batch trials for Ax

[schmoelder]

Running batch trials is currently disabled in favor of early stopping, as both strategies were conflicting on Ax' side. Manually performing early stopping once the best-point/best-pareto can be returned from Ax.

• Improve obtaining best-point/best-pareto during optimization
• Improve setting termination criteria. For this it would be useful, if Ax could return the best point. This is currently not implemented
• Adjust the number of cores to the number of batch trials, once batch mode is enabled again

Starting Point

This is a classmethod of AxInterface

from ax.modelbridge.modelbridge_utils import (
    observed_pareto_frontier as observed_pareto,
    predicted_pareto_frontier as predicted_pareto,
)

def _get_best_results(self, modelbridge, optimization_config):
    raise NotImplementedError("does not work so far.")
    n_obj = self.optimization_problem.n_objectives 
    if n_obj > 1:
        pareto_optimal_observations = observed_pareto(
            modelbridge, 
            optimization_config=optimization_config
        )
        pareto_optimal_predictions = predicted_pareto(
            modelbridge, 
            optimization_config=optimization_config
        )

        # here also work needs to be done.

    else:
        arm, obj = modelbridge.model_best_point()
        x = np.array(list(arm.parameters.values()))
        metrics = optimization_config.objective.metric_names
        assert len(metrics) == 1, "In single objective problems, only 1 Metric is allowed."
        f = obj[0][metrics[0]]

        return np.array([x], ndmin=2), np.array([f], ndmin=2)

then in the optimization loop:

X_opt, F_opt = self._get_best_results(
    modelbridge=modelbridge, 
    optimization_config=optimization_config
)

However, currently Ax does not return the best point, the reason for this is unclear

---

Here is another removed old snipped from code. May be another starting point for providing the current optimum to self._post_processing in the training loop.

self._data = ax.Data.from_multiple_data([self._data, trial.fetch_data()])
new_value = trial.fetch_data().df["mean"].min()

print(
    f"Iteration: Best in iteration {new_value:.3f}, "
    f"Best so far: {self._data.df['mean'].min():.3f}"
)

[flo-schu]

General Ax usage

To tackle this problem, I will first have a look at this multi-objective tutorial written for the developer API (https://ax.dev/tutorials/saasbo_nehvi.html) and check the existing implementation of Ax to see if we can use things closer to what the developers have in mind.

Also pytorch's tutorial seems like a good place to start to understand Ax better: https://pytorch.org/tutorials/intermediate/ax_multiobjective_nas_tutorial.html

Early stopping

There is a trial level early stopping tutorial, which is interesting, because it can cancel jobs while running, when they are "not promising" https://ax.dev/tutorials/early_stopping/early_stopping.html

Global level early stopping is definitely relevant this: https://ax.dev/tutorials/gss.html. Point 3 in the tutorial "Write your own custom global strategy" seems to be highly relevant.

[flo-schu]

I have raised an issue regarding the implementation of global level early stopping for batch trials facebook/Ax#2677

Maybe we will get some leads for our implementation.

[flo-schu]

facebook/Ax@1fa346b implements Batch trial support