feat: Update Surrogate a final time after optimization process finished (#132)

Author: sumny

R/OptimizerMbo.R

@@ -15,6 +15,13 @@
 #'
 #' Termination is handled via a [bbotk::Terminator] part of the [bbotk::OptimInstance] to be optimized.
 #'
+#' Note that in general the [Surrogate] is updated one final time on all available data after the optimization process has terminated.
+#' However, in certain scenarios this is not always possible or meaningful, e.g., when using [bayesopt_parego()] for multi-objective optimization
+#' which uses a surrogate that relies on a scalarization of the objectives.
+#' It is therefore recommended to manually inspect the [Surrogate] after optimization if it is to be used, e.g., for visualization purposes to make
+#' sure that it has been properly updated on all available data.
+#' If this final update of the [Surrogate] could not be performed successfully, a warning will be logged.
+#'
 #' @section Archive:
 #' The [bbotk::Archive] holds the following additional columns that are specific to MBO algorithms:
 #'   * `[acq_function$id]` (`numeric(1)`)\cr
@@ -274,6 +281,7 @@ OptimizerMbo = R6Class("OptimizerMbo",
 
     .optimize = function(inst) {
       # FIXME: this needs more checks for edge cases like eips or loop_function bayesopt_parego then default_surrogate should use one learner
+
       if (is.null(self$loop_function)) {
         self$loop_function = default_loop_function(inst)
       }
@@ -303,6 +311,17 @@ OptimizerMbo = R6Class("OptimizerMbo",
       check_packages_installed(self$packages, msg = sprintf("Package '%%s' required but not installed for Optimizer '%s'", format(self)))
 
       invoke(self$loop_function, instance = inst, surrogate = self$surrogate, acq_function = self$acq_function, acq_optimizer = self$acq_optimizer, .args = self$args)
+
+      on.exit({
+        tryCatch(
+          {
+            self$surrogate$update()
+          }, surrogate_update_error = function(error_condition) {
+            logger = lgr::get_logger("bbotk")
+            logger$warn("Could not update the surrogate a final time after the optimization process has terminated.")
+          }
+        )
+      })
     },
 
     .assign_result = function(inst) {

man/mlr_optimizers_mbo.Rd

@@ -3,7 +3,7 @@ test_that("OptimizerMbo works", {
   skip_if_not_installed("mlr3learners")
   skip_if_not_installed("DiceKriging")
   skip_if_not_installed("rgenoud")
-  
+
   optimizer = OptimizerMbo$new()
   expect_r6(optimizer, classes = "OptimizerMbo")
 
@@ -199,3 +199,22 @@ test_that("OptimizerMbo reset", {
   expect_r6(optimizer$acq_optimizer, "AcqOptimizer")
 })
 
+test_that("OptimizerMbo up to date surrogate after optimization", {
+  skip_if_not_installed("mlr3learners")
+  skip_if_not_installed("DiceKriging")
+
+  surrogate = srlrn(lrn("regr.km", covtype = "matern5_2", optim.method = "gen", control = list(trace = FALSE), nugget.stability = 10^-8))
+  acq_optimizer = acqo(opt("random_search", batch_size = 2L), terminator = trm("evals", n_evals = 2L))
+  optimizer = opt("mbo", surrogate = surrogate, acq_optimizer = acq_optimizer)
+  instance = MAKE_INST_1D(terminator = trm("evals", n_evals = 5L))
+  optimizer$optimize(instance)
+
+  expect_equal(surrogate, optimizer$surrogate)
+
+  expect_true(surrogate$learner$state$train_task$nrow == 5L)
+
+  predictions = surrogate$predict(instance$archive$data[, instance$archive$cols_x, with = FALSE])
+  expect_true(all(sqrt((predictions$mean - instance$archive$data$y) ^ 2) < 1e-4))
+  expect_true(all(predictions$se < 1e-4))
+})
+