Implementation of Ax Optimizer

.github/workflows/r-cmd-check.yml

@@ -54,11 +54,32 @@ jobs:
         with:
           r-version: ${{ matrix.config.r }}
 
+      - uses: actions/setup-python@v5
+        with:
+          python-version: '3.11'
+
       - uses: r-lib/actions/setup-r-dependencies@v2
         with:
           extra-packages: any::rcmdcheck
           needs: check
 
+      - name: Create Python virtual environments
+        run: |
+          python3 -m venv ~/.virtualenvs/r-hebo
+          ~/.virtualenvs/r-hebo/bin/pip install --upgrade pip --quiet
+          ~/.virtualenvs/r-hebo/bin/pip install torch --index-url https://download.pytorch.org/whl/cpu --quiet
+          ~/.virtualenvs/r-hebo/bin/pip install hebo --quiet
+          python3 -m venv ~/.virtualenvs/r-smac
+          ~/.virtualenvs/r-smac/bin/pip install --upgrade pip --quiet
+          ~/.virtualenvs/r-smac/bin/pip install smac ConfigSpace --quiet
+          python3 -m venv ~/.virtualenvs/r-optuna
+          ~/.virtualenvs/r-optuna/bin/pip install --upgrade pip --quiet
+          ~/.virtualenvs/r-optuna/bin/pip install optuna --quiet
+          python3 -m venv ~/.virtualenvs/r-ax
+          ~/.virtualenvs/r-ax/bin/pip install --upgrade pip --quiet
+          ~/.virtualenvs/r-ax/bin/pip install torch --index-url https://download.pytorch.org/whl/cpu --quiet
+          ~/.virtualenvs/r-ax/bin/pip install ax-platform --quiet
+
       - uses: mxschmitt/action-tmate@v3
         if: ${{ github.event_name == 'workflow_dispatch' && inputs.debug_enabled }}
         with:

DESCRIPTION

@@ -11,7 +11,8 @@ Authors@R: c(
     person("Bernd", "Bischl", , "bernd_bischl@gmx.net", role = "aut",
            comment = c(ORCID = "0000-0001-6002-6980")),
     person("Martin", "Binder", , "martin.binder@mail.com", role = "aut"),
-    person("Olaf", "Mersmann", , "olafm@statistik.tu-dortmund.de", role = "ctb")
+    person("Olaf", "Mersmann", , "olafm@statistik.tu-dortmund.de", role = "ctb"),
+    person("Alexander", "Winterstetter", "alexanderwinterstetter@gmail.de", role = "ctb")
   )
 Description: Features highly configurable search spaces via the 'paradox'
     package and optimizes every user-defined objective function. The
@@ -35,18 +36,24 @@ Imports:
     R6
 Suggests:
     adagio,
+    callr,
     emoa,
     GenSA,
+    reticulate,
     irace (>= 4.0.0),
     knitr,
     mirai,
     nloptr,
     processx,
     progressr,
     redux,
+    reticulate,
     RhpcBLASctl,
     rush (>= 1.0.0),
-    testthat (>= 3.0.0)
+    testthat (>= 3.0.0),
+    withr
+Remotes:
+    mlr-org/paradox
 Config/testthat/edition: 3
 Config/testthat/parallel: false
 Encoding: UTF-8
@@ -85,16 +92,20 @@ Collate:
     'OptimizerAsyncGridSearch.R'
     'OptimizerAsyncRandomSearch.R'
     'OptimizerBatch.R'
+    'OptimizerBatchAx.R'
     'OptimizerBatchChain.R'
     'OptimizerBatchCmaes.R'
     'OptimizerBatchDesignPoints.R'
     'OptimizerBatchFocusSearch.R'
     'OptimizerBatchGenSA.R'
     'OptimizerBatchGridSearch.R'
+    'OptimizerBatchHEBO.R'
     'OptimizerBatchIrace.R'
     'OptimizerBatchLocalSearch.R'
     'OptimizerBatchNLoptr.R'
+    'OptimizerBatchOptuna.R'
     'OptimizerBatchRandomSearch.R'
+    'OptimizerBatchSmac3.R'
     'Progressor.R'
     'mlr_terminators.R'
     'Terminator.R'
@@ -114,6 +125,9 @@ Collate:
     'bibentries.R'
     'conditions.R'
     'helper.R'
+    'helper_ax.R'
+    'helper_hebo.R'
+    'helper_optuna.R'
     'local_search.R'
     'mlr_callbacks.R'
     'mlr_test_functions.R'

NAMESPACE

@@ -42,16 +42,20 @@ export(OptimizerAsyncDesignPoints)
 export(OptimizerAsyncGridSearch)
 export(OptimizerAsyncRandomSearch)
 export(OptimizerBatch)
+export(OptimizerBatchAx)
 export(OptimizerBatchChain)
 export(OptimizerBatchCmaes)
 export(OptimizerBatchDesignPoints)
 export(OptimizerBatchFocusSearch)
 export(OptimizerBatchGenSA)
 export(OptimizerBatchGridSearch)
+export(OptimizerBatchHEBO)
 export(OptimizerBatchIrace)
 export(OptimizerBatchLocalSearch)
 export(OptimizerBatchNLoptr)
+export(OptimizerBatchOptuna)
 export(OptimizerBatchRandomSearch)
+export(OptimizerBatchSmac3)
 export(Terminator)
 export(TerminatorClockTime)
 export(TerminatorCombo)

R/OptimizerAsync.R

@@ -20,13 +20,13 @@
 #' There are three types of workers:
 #'
 #' - "mirai": Workers are started with \CRANpkg{mirai} on local or remote machines.
-#'   See `$start_workers()` in [Rush] for more details.
+#'   See `$start_workers()` in [rush::Rush] for more details.
 #'   [mirai::daemons()] must be created before starting the optimization.
 #' - "processx": Workers are started as local processes with \CRANpkg{processx}.
 
-#'   See `$start_local_workers()` in [Rush] for more details.
+#'   See `$start_local_workers()` in [rush::Rush] for more details.
 #' - "script": Workers are started by the user with a custom script.
-#'   See `$create_worker_script()` in [Rush] for more details.
+#'   See `$create_worker_script()` in [rush::Rush] for more details.
 #'
 #' The workers are started when the `$optimize()` method is called.
 #' The main process waits until at least one worker is running.

R/OptimizerBatchAx.R

@@ -0,0 +1,129 @@
+#' @title Optimization via Ax (Adaptive Experimentation Platform)
+#'
+#' @include Optimizer.R
+#' @name mlr_optimizers_ax
+#'
+#' @description
+#' `OptimizerBatchAx` class that implements optimization via the Python
+#' Ax framework. Calls Ax via \CRANpkg{reticulate}.
+#'
+#' @templateVar id ax
+#' @template section_dictionary_optimizers
+#'
+#' @section Parameters:
+#' \describe{
+#' \item{`n_sobol`}{`integer(1)`\cr
+#'   Number of quasi-random (Sobol) initialization trials before Bayesian optimization starts.
+#'   If `NULL` (default), Ax determines the number automatically based on the search space.}
+#' \item{`seed`}{`integer(1)`\cr
+#'   Random seed passed to the `AxClient` for reproducibility.
+#'   Default is `NULL` (no fixed seed).}
+#' }
+#'
+#' @template section_progress_bars
+#'
+#' @export
+OptimizerBatchAx = R6Class(
+  "OptimizerBatchAx",
+  inherit = OptimizerBatch,
+  public = list(
+    #' @description
+    #' Creates a new instance of this [R6][R6::R6Class] class.
+    initialize = function() {
+      param_set = ps(
+        n_sobol = p_int(lower = 1L, special_vals = list(NULL)),
+        seed = p_int(lower = 1L, special_vals = list(NULL))
+      )
+
+      super$initialize(
+        id = "ax",
+        param_set = param_set,
+        param_classes = c("ParamDbl", "ParamInt", "ParamFct", "ParamLgl"),
+        properties = c("single-crit", "multi-crit"),
+        packages = "reticulate",
+        label = "Ax",
+        man = "bbotk::mlr_optimizers_ax"
+      )
+    }
+  ),
+
+  private = list(
+    .optimize = function(inst) {
+      assert_python_packages("ax")
+      ax_service = reticulate::import("ax.service.ax_client")
+
+      pv = self$param_set$values
+      search_space = inst$search_space
+      codomain = inst$archive$codomain
+
+      # Build Ax parameter definitions from paradox search space
+      parameters = paramset_to_ax_params(search_space)
+
+      # Build objectives dict from codomain
+      objective_ids = codomain$ids()
+      objectives = setNames(
+        lapply(objective_ids, function(id) {
+          minimize = !("maximize" %in% codomain$params[[id]]$tags)
+          ax_service$ObjectiveProperties(minimize = minimize)
+        }),
+        objective_ids
+      )
+
+      # Build AxClient args; optionally set seed and custom generation strategy
+      client_args = list(verbose_logging = FALSE)
+      if (!is.null(pv$seed)) {
+        client_args$random_seed = as.integer(pv$seed)
+      }
+      if (!is.null(pv$n_sobol)) {
+        ax_gs = reticulate::import("ax.modelbridge.generation_strategy")
+        ax_models = reticulate::import("ax.modelbridge.registry")
+        client_args$generation_strategy = ax_gs$GenerationStrategy(
+          steps = list(
+            # initial data points
+            ax_gs$GenerationStep(
+              model = ax_models$Models$SOBOL,
+              num_trials = pv$n_sobol
+            ),
+
+            ax_gs$GenerationStep(
+              model = ax_models$Models$BOTORCH_MODULAR,
+              num_trials = -1L
+            )
+          )
+        )
+      }
+
+      ax_client = invoke(ax_service$AxClient, .args = client_args)
+
+      ax_client$create_experiment(
+        name = "bbotk_experiment",
+        parameters = parameters,
+        objectives = objectives
+      )
+
+      # Ask-and-tell loop
+      repeat {
+        result = tryCatch(
+          ax_client$get_next_trial(),
+          error = function(e) NULL
+        )
+        if (is.null(result)) {
+          break
+        }
+
+        params = reticulate::py_to_r(result[[1L]])
+        trial_index = result[[2L]]
+
+        xdt = ax_params_to_xdt(params, search_space)
+        inst$eval_batch(xdt)
+
+        n = nrow(inst$archive$data)
+        y_vals = unlist(inst$archive$data[n, inst$archive$cols_y, with = FALSE])
+
+        ax_client$complete_trial(trial_index = trial_index, raw_data = as.list(y_vals))
+      }
+    }
+  )
+)
+
+mlr_optimizers$add("ax", OptimizerBatchAx)