Repository: optuna/optuna
Pull request: Add `logei_candidate_func` and make it default when available
## Motivation
Resolve #4546.

The rationale for this improvement and the official benchmarks are available [here](https://github.com/pytorch/botorch/pull/1565).

## Description of the changes
* Add `logei_candidate_func`.
* Use `logei_candidate_func` as default when there is no constraint. ~and `LogExpectedImprovement` is available.~
* Modify tests accordingly.

Note that `LogExpectedImprovement` has the following limitations:
* Constraints are not available; supporting constraints needs quasi-Monte Carlo sampling.
* `LogExpectedImprovement` is only available in `botorch >=0.8.1`.

In those cases, this PR simply falls back to `qei_candidate_func`.



Problem statement:
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.

Hints from issue discussion:

Commit message: Update optuna/integration/botorch.py

Co-authored-by: Naoto Mizuno <naotomizuno@preferred.jp>

Patch to review:
diff --git a/docs/source/reference/integration.rst b/docs/source/reference/integration.rst
index b4288dcb621..4ab72a0c02e 100644
--- a/docs/source/reference/integration.rst
+++ b/docs/source/reference/integration.rst
@@ -26,6 +26,7 @@ BoTorch
    :nosignatures:
 
    optuna.integration.BoTorchSampler
+   optuna.integration.botorch.logei_candidates_func
    optuna.integration.botorch.qei_candidates_func
    optuna.integration.botorch.qehvi_candidates_func
    optuna.integration.botorch.qnehvi_candidates_func
diff --git a/optuna/integration/botorch.py b/optuna/integration/botorch.py
index 15d9c31106b..26943af0103 100644
--- a/optuna/integration/botorch.py
+++ b/optuna/integration/botorch.py
@@ -63,6 +63,100 @@ 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,
+        sampler=_get_sobol_qmc_normal_sampler(256),
+    )
+
+    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 +169,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 +513,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 +529,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")
@@ -522,7 +620,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 +752,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]
diff --git a/tests/integration_tests/test_botorch.py b/tests/integration_tests/test_botorch.py
index 1290e483da1..9f157024e8a 100644
--- a/tests/integration_tests/test_botorch.py
+++ b/tests/integration_tests/test_botorch.py
@@ -1,3 +1,4 @@
+import sys
 from typing import Any
 from typing import Optional
 from typing import Sequence
@@ -32,6 +33,9 @@
 
 @pytest.mark.parametrize("n_objectives", [1, 2, 4])
 def test_botorch_candidates_func_none(n_objectives: int) -> None:
+    if n_objectives == 1 and sys.version_info < (3, 8):
+        pytest.skip("Python >=3.8 is required for logei_candidates_func.")
+
     n_trials = 3
     n_startup_trials = 2
 
@@ -46,7 +50,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 +93,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 +102,13 @@ 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 sys.version_info < (
+        3,
+        8,
+        0,
+    ):
+        pytest.skip("LogExpectedImprovement is not available in Python <3.8.")
+
     n_trials = 4
     n_startup_trials = 2
 
diff --git a/tests/samplers_tests/test_samplers.py b/tests/samplers_tests/test_samplers.py
index 929ff65f846..5212bed120f 100644
--- 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,
         ),
     ],


Code review comments:
diff_hunk: @@ -63,6 +63,100 @@ 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
path: optuna/integration/botorch.py
text: @user1:
A nit.
```suggestion
            Objective constraints. This option is not supported in ``logei_candidates_func`` and
```

diff_hunk: @@ -63,6 +63,100 @@ 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,
+        sampler=_get_sobol_qmc_normal_sampler(256),
path: optuna/integration/botorch.py
text: @user1:
I cannot find the `sampler` argument in the `LogExpectedImprovement` reference. Since it employs an analytical approach according to the module name, it might not require it.

```suggestion
```

@author:
Nice catch! Thaks for the comment.

diff_hunk: @@ -32,6 +33,9 @@
 
 @pytest.mark.parametrize("n_objectives", [1, 2, 4])
 def test_botorch_candidates_func_none(n_objectives: int) -> None:
+    if n_objectives == 1 and sys.version_info < (3, 8):
path: tests/integration_tests/test_botorch.py
text: @user1:
The logic used to bypass the tests seems somewhat indirect. It could be clearer if it directly verifies the version of botorch, in a similar manner to what's demonstrated in `optuna/integration/botorch.py`. Would that be a feasible approach for you to consider?

https://github.com/optuna/optuna/blob/2b6addc365003eaab9e20bacb964ff15ed057da7/optuna/integration/botorch.py#L42


@author:
```suggestion
    if n_objectives == 1 and version.parse(botorch.version.version) < version.parse("0.8.1"):
```
Thanks. I'll make the change.