Repository: optuna/optuna

Issue:
Use `LogExpectedImprovement` for `BoTorchSampler`
Currently, `BoTorchSampler` uses `qExpectedImprovement` acquisition function. This relies on quasi-Monte Carlo sampling to compute the acquisition function, which is not accurate when the probability of improvement is very low. 

Quasi-Monte Carlo methods are last-resort methods for computing acquisition functions that are not analytically computable. This is the case when `q >= 2`, i.e. in batch optimization, but Optuna does not do batch optimization, so the expected improvement can be computed analytically.

BoTorch has a better acquisition function class `botorch.acquisition.analytic.LogExpectedImprovement` (defined as `log` of  `ExpectedImprovement`). Maximizing this function is mathematically equivalent to maximizing `ExpectedImprovement`, but numerically maximizing `ExpectedImprovement` is difficult because it tends to be very small for the vast area of the search space (and also susceptible to precision loss). Instead, `LogExpectedImprovement` can provide more meaningful gradient information for the optimizer.



Discussion:

Resolving patch:
diff --git a/optuna/integration/botorch.py b/optuna/integration/botorch.py
--- a/optuna/integration/botorch.py
+++ b/optuna/integration/botorch.py
@@ -63,6 +63,99 @@ def _get_sobol_qmc_normal_sampler(num_samples: int) -> SobolQMCNormalSampler:
 
 _logger = logging.get_logger(__name__)
 
+with try_import() as _imports_logei:
+    from botorch.acquisition.analytic import LogExpectedImprovement
+
+
+@experimental_func("3.3.0")
+def logei_candidates_func(
+    train_x: "torch.Tensor",
+    train_obj: "torch.Tensor",
+    train_con: Optional["torch.Tensor"],
+    bounds: "torch.Tensor",
+    pending_x: Optional["torch.Tensor"],
+) -> "torch.Tensor":
+    """Log Expected Improvement (LogEI).
+
+    The default value of ``candidates_func`` in :class:`~optuna.integration.BoTorchSampler`
+    with single-objective optimization for non-constrained problems.
+
+    Args:
+        train_x:
+            Previous parameter configurations. A ``torch.Tensor`` of shape
+            ``(n_trials, n_params)``. ``n_trials`` is the number of already observed trials
+            and ``n_params`` is the number of parameters. ``n_params`` may be larger than the
+            actual number of parameters if categorical parameters are included in the search
+            space, since these parameters are one-hot encoded.
+            Values are not normalized.
+        train_obj:
+            Previously observed objectives. A ``torch.Tensor`` of shape
+            ``(n_trials, n_objectives)``. ``n_trials`` is identical to that of ``train_x``.
+            ``n_objectives`` is the number of objectives. Observations are not normalized.
+        train_con:
+            Objective constraints. This option is not supported in ``logei_candidates_func`` and
+            must be :obj:`None`.
+        bounds:
+            Search space bounds. A ``torch.Tensor`` of shape ``(2, n_params)``. ``n_params`` is
+            identical to that of ``train_x``. The first and the second rows correspond to the
+            lower and upper bounds for each parameter respectively.
+        pending_x:
+            Pending parameter configurations. A ``torch.Tensor`` of shape
+            ``(n_pending, n_params)``. ``n_pending`` is the number of the trials which are already
+            suggested all their parameters but have not completed their evaluation, and
+            ``n_params`` is identical to that of ``train_x``.
+
+    Returns:
+        Next set of candidates. Usually the return value of BoTorch's ``optimize_acqf``.
+
+    """
+
+    # We need botorch >=0.8.1 for LogExpectedImprovement.
+    if not _imports_logei.is_successful():
+        raise ImportError(
+            "logei_candidates_func requires botorch >=0.8.1. "
+            "Please upgrade botorch or use qei_candidates_func as candidates_func instead."
+        )
+
+    if train_obj.size(-1) != 1:
+        raise ValueError("Objective may only contain single values with logEI.")
+    if train_con is not None:
+        raise ValueError(
+            "Constraint is not supported with logei_candidates_func. "
+            + "Please use qei_candidates_func instead."
+        )
+    else:
+        train_y = train_obj
+        best_f = train_obj.max()
+
+    train_x = normalize(train_x, bounds=bounds)
+
+    model = SingleTaskGP(train_x, train_y, outcome_transform=Standardize(m=train_y.size(-1)))
+    mll = ExactMarginalLogLikelihood(model.likelihood, model)
+    fit_gpytorch_mll(mll)
+
+    acqf = LogExpectedImprovement(
+        model=model,
+        best_f=best_f,
+    )
+
+    standard_bounds = torch.zeros_like(bounds)
+    standard_bounds[1] = 1
+
+    candidates, _ = optimize_acqf(
+        acq_function=acqf,
+        bounds=standard_bounds,
+        q=1,
+        num_restarts=10,
+        raw_samples=512,
+        options={"batch_limit": 5, "maxiter": 200},
+        sequential=True,
+    )
+
+    candidates = unnormalize(candidates.detach(), bounds=bounds)
+
+    return candidates
+
 
 @experimental_func("2.4.0")
 def qei_candidates_func(
@@ -75,7 +168,7 @@ def qei_candidates_func(
     """Quasi MC-based batch Expected Improvement (qEI).
 
     The default value of ``candidates_func`` in :class:`~optuna.integration.BoTorchSampler`
-    with single-objective optimization.
+    with single-objective optimization for constrained problems.
 
     Args:
         train_x:
@@ -419,6 +512,8 @@ def qparego_candidates_func(
 
 def _get_default_candidates_func(
     n_objectives: int,
+    has_constraint: bool,
+    consider_running_trials: bool,
 ) -> Callable[
     [
         "torch.Tensor",
@@ -433,8 +528,10 @@ def _get_default_candidates_func(
         return qparego_candidates_func
     elif n_objectives > 1:
         return qehvi_candidates_func
-    else:
+    elif has_constraint or consider_running_trials:
         return qei_candidates_func
+    else:
+        return logei_candidates_func
 
 
 @experimental_class("2.4.0")
@@ -467,8 +564,11 @@ class BoTorchSampler(BaseSampler):
             :obj:`None`. For any constraints that failed to compute, the tensor will contain
             NaN.
 
-            If omitted, it is determined automatically based on the number of objectives. If the
-            number of objectives is one, Quasi MC-based batch Expected Improvement (qEI) is used.
+            If omitted, it is determined automatically based on the number of objectives and
+            whether a constraint is specified. If the
+            number of objectives is one and no constraint is specified, log-Expected Improvement
+            is used. If constraints are specified, quasi MC-based batch Expected Improvement
+            (qEI) is used.
             If the number of objectives is either two or three, Quasi MC-based
             batch Expected Hypervolume Improvement (qEHVI) is used. Otherwise, for larger number
             of objectives, the faster Quasi MC-based extended ParEGO (qParEGO) is used.
@@ -522,7 +622,7 @@ def __init__(
         ] = None,
         constraints_func: Optional[Callable[[FrozenTrial], Sequence[float]]] = None,
         n_startup_trials: int = 10,
-        consider_running_trials: bool = True,
+        consider_running_trials: bool = False,
         independent_sampler: Optional[BaseSampler] = None,
         seed: Optional[int] = None,
         device: Optional["torch.device"] = None,
@@ -654,7 +754,11 @@ def sample_relative(
         bounds.transpose_(0, 1)
 
         if self._candidates_func is None:
-            self._candidates_func = _get_default_candidates_func(n_objectives=n_objectives)
+            self._candidates_func = _get_default_candidates_func(
+                n_objectives=n_objectives,
+                has_constraint=con is not None,
+                consider_running_trials=self._consider_running_trials,
+            )
 
         completed_values = values[:n_completed_trials]
         completed_params = params[:n_completed_trials]


Test patch:
diff --git a/tests/integration_tests/test_botorch.py b/tests/integration_tests/test_botorch.py
--- a/tests/integration_tests/test_botorch.py
+++ b/tests/integration_tests/test_botorch.py
@@ -5,6 +5,8 @@
 from unittest.mock import patch
 import warnings
 
+import botorch
+from packaging import version
 import pytest
 
 import optuna
@@ -32,6 +34,9 @@
 
 @pytest.mark.parametrize("n_objectives", [1, 2, 4])
 def test_botorch_candidates_func_none(n_objectives: int) -> None:
+    if n_objectives == 1 and version.parse(botorch.version.version) < version.parse("0.8.1"):
+        pytest.skip("botorch >=0.8.1 is required for logei_candidates_func.")
+
     n_trials = 3
     n_startup_trials = 2
 
@@ -46,7 +51,7 @@ def test_botorch_candidates_func_none(n_objectives: int) -> None:
 
     # TODO(hvy): Do not check for the correct candidates function using private APIs.
     if n_objectives == 1:
-        assert sampler._candidates_func is integration.botorch.qei_candidates_func
+        assert sampler._candidates_func is integration.botorch.logei_candidates_func
     elif n_objectives == 2:
         assert sampler._candidates_func is integration.botorch.qehvi_candidates_func
     elif n_objectives == 4:
@@ -89,6 +94,7 @@ def candidates_func(
 @pytest.mark.parametrize(
     "candidates_func, n_objectives",
     [
+        (integration.botorch.logei_candidates_func, 1),
         (integration.botorch.qei_candidates_func, 1),
         (integration.botorch.qehvi_candidates_func, 2),
         (integration.botorch.qparego_candidates_func, 4),
@@ -97,6 +103,11 @@ def candidates_func(
     ],
 )
 def test_botorch_specify_candidates_func(candidates_func: Any, n_objectives: int) -> None:
+    if candidates_func == integration.botorch.logei_candidates_func and version.parse(
+        botorch.version.version
+    ) < version.parse("0.8.1"):
+        pytest.skip("LogExpectedImprovement is not available in botorch <0.8.1.")
+
     n_trials = 4
     n_startup_trials = 2
 
diff --git a/tests/samplers_tests/test_samplers.py b/tests/samplers_tests/test_samplers.py
--- a/tests/samplers_tests/test_samplers.py
+++ b/tests/samplers_tests/test_samplers.py
@@ -3,6 +3,7 @@
 from multiprocessing.managers import DictProxy
 import os
 import pickle
+import sys
 from typing import Any
 from typing import Callable
 from typing import Dict
@@ -25,6 +26,8 @@
 from optuna.distributions import CategoricalDistribution
 from optuna.distributions import FloatDistribution
 from optuna.distributions import IntDistribution
+from optuna.integration.botorch import logei_candidates_func
+from optuna.integration.botorch import qei_candidates_func
 from optuna.samplers import BaseSampler
 from optuna.study import Study
 from optuna.testing.objectives import fail_objective
@@ -56,7 +59,20 @@
         optuna.samplers.NSGAIIISampler,
         optuna.samplers.QMCSampler,
         pytest.param(
-            lambda: optuna.integration.BoTorchSampler(n_startup_trials=0),
+            lambda: optuna.integration.BoTorchSampler(
+                n_startup_trials=0,
+                candidates_func=logei_candidates_func,
+            ),
+            marks=[
+                pytest.mark.integration,
+                pytest.mark.skipif(sys.version_info < (3, 8), reason="Requires Python 3.8+"),
+            ],
+        ),
+        pytest.param(
+            lambda: optuna.integration.BoTorchSampler(
+                n_startup_trials=0,
+                candidates_func=qei_candidates_func,
+            ),
             marks=pytest.mark.integration,
         ),
     ],