Add qConstrainedLogEI

[sawa3030]

Motivation

PR #6640 introduced a Monte Carlo-based acquisition function for unconstrained single-objective optimization in GPSampler. This PR extends the approach to constrained single-objective optimization.

Description of the changes

• Add qConstrainedLogEI.
• Use qConstrainedLogEI as the default acquisition function for constrained single-objective optimization.

[sawa3030]

Benchmark

I compared the master branch with this PR by simulating parallel optimization using several values of tau.

The optimization performance was evaluated in terms of the number of completed trials. Overall, this PR improves optimization performance compared with the master branch. However, when the feasible region is small, a relatively large tau overly smooths the constraint indicator, making it more difficult to identify feasible points and consequently degrading optimization performance.

Based on the results shown in the graph below, I selected tau=0.01, which provides a good balance between numerical smoothness and accurate feasibility estimation.

The benchmark uses the following two synthetic problems from Bayesian Optimization with Inequality Constraints (Gardner et al., ICML 2014):

• Problem 1
  • Objective: cos(2x) cos(y) + sin(x)
  • Constraint: cos(x) cos(y) - sin(x) sin(y) <= 0.5
• Problem 2
  • Objective: sin(x) + y
  • Constraint: sin(x) sin(y) <= -0.95

Simulation Code

from __future__ import annotations

import argparse
import itertools
import pickle
import time
from pathlib import Path
from typing import Any

import optuna
import optunahub
from optuna.distributions import (
    BaseDistribution,
    CategoricalDistribution,
    FloatDistribution,
    IntDistribution,
)
import numpy as np

def simulate(
    n_workers: int,
    n_trials: int,
    n_startup_trials: int,
    seed: int,
    dataset_id: int,
    tau: float,
    use_qmc: bool,
) -> optuna.Study:
    if n_workers <= 0:
        raise ValueError("n_workers must be >= 1")

    def objective(x: float, y: float) -> float:
        # return float(np.cos(2*x) * np.cos(y) + np.sin(x))
        return float(np.sin(x) + y)

    def constraints(trial: optuna.trial.FrozenTrial) -> tuple[float]:
        x = trial.params["x"]
        y = trial.params["y"]
        # c = float(np.cos(x) * np.cos(y) - np.sin(x) * np.sin(y) - 0.5)
        c = float(np.sin(x)*np.sin(y) + 0.95)
        return (c,)
        
    def feasible(trial: optuna.trial.FrozenTrial) -> bool:
        return all(c <= 0 for c in constraints(trial))

    sampler = optuna.samplers.GPSampler(
        n_startup_trials=n_startup_trials,
        seed=seed,
        constraints_func=constraints,
    )
    sampler._tau = tau
    sampler._use_qmc = use_qmc
    study = optuna.create_study(sampler=sampler)
    start_time = time.perf_counter()

    pending: list[tuple[optuna.Trial, dict[str, Any]] | None] = [None] * n_workers
    n_suggested = 0
    n_completed = 0

    for worker_id in itertools.cycle(range(n_workers)):
        if n_completed >= n_trials:
            break

        if pending[worker_id] is not None:
            previous_trial, previous_params = pending[worker_id]
            value = objective(**previous_params)
            study.tell(previous_trial, value)
            pending[worker_id] = None
            n_completed += 1
            if n_completed >= n_trials:
                break

        if n_suggested < n_trials:
            trial = study.ask()
            x = trial.suggest_float("x", 0.0, 2 * np.pi)
            y = trial.suggest_float("y", 0.0, 2 * np.pi)
            params = {"x": x, "y": y}
            trial.set_user_attr("trial_num", trial._trial_id + 1)  # Simulate cumulative time as trial number + 1
            pending[worker_id] = (trial, params)
            n_suggested += 1

    trials = [
        t
        for t in study.trials[n_startup_trials+n_workers-1:]
        if t.state == optuna.trial.TrialState.COMPLETE and feasible(t)
    ]
    trials = trials[: max(0, n_trials - n_startup_trials - n_workers + 1)]

    new_study = optuna.create_study()
    new_study.add_trials(trials)
    return new_study


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--label", required=True)
    parser.add_argument("--out-dir", default="./gp_simulator_results_without_evaltime_n5_bbo_dataset10")
    parser.add_argument("--n-workers", type=int, default=5)
    parser.add_argument("--n-trials", type=int, default=100)
    parser.add_argument("--n-startup-trials", type=int, default=10)
    parser.add_argument("--n-seeds", type=int, default=10)
    parser.add_argument("--dataset-id", type=int, default=10)
    parser.add_argument("--tau", type=float, default=0.01)
    parser.add_argument("--use_qmc", type=bool, default=False)

    study_list: list[optuna.Study] = []
    for seed in range(args.n_seeds):
        study = simulate(
            n_workers=args.n_workers,
            n_trials=args.n_trials,
            n_startup_trials=args.n_startup_trials,
            seed=seed,
            dataset_id=args.dataset_id,
            tau=args.tau,
            use_qmc=args.use_qmc,
        )
        study_list.append(study)

    out_dir = Path(args.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    out_path = out_dir / f"{args.label}.pickle"

    with out_path.open("wb") as f:
        pickle.dump(study_list, f)

    print(f"Saved {len(study_list)} studies to {out_path}")


if __name__ == "__main__":
    main()

Plot Code

from __future__ import annotations

import argparse
from pathlib import Path
import pickle
import re

import matplotlib.pyplot as plt
import optuna
import optunahub
import random


plot_target_over_time = optunahub.load_module(
    "visualization/plot_target_over_time"
).plot_target_over_time


def load_study_list(result_dir: Path, label: str) -> list[optuna.Study]:
    path = result_dir / f"{label}.pickle"
    with path.open("rb") as f:
        study_list = pickle.load(f)
    return study_list


def trim_to_same_length(study_list: list[optuna.Study]) -> list[optuna.Study]:
    min_len = min(len(study.trials) for study in study_list)

    trimmed = []
    for study in study_list:
        new_study = optuna.create_study(directions=study.directions)
        new_study.add_trials(study.trials[:min_len])
        trimmed.append(new_study)

    return trimmed


def get_style(label: str) -> dict[str, str]:
    if label == "master":
        return {
            "color": "blue",
            "marker": "s",
            "ls": "dashed",
            "plot_label": "master",
        }
    elif label == "qConstrainedLogEI":
        return {
            "color": "green",
            "marker": "D",
            "ls": "dashdot",
            # "plot_label": "qLogEI (n_qmc_samples=128)",
            "plot_label": "qConstrainedLogEI",
        }
    elif label == "qConstrainedLogEI-tau0-01":
        return {
            "color": "red",
            "marker": "P",
            "ls": "dashdot",
            "plot_label": "qConstrainedLogEI (tau=0.01)",
        }
    elif label == "qConstrainedLogEI-tau0-1":
        return {
            "color": "purple",
            "marker": "X",
            "ls": "dashdot",
            "plot_label": "qConstrainedLogEI (tau=0.1)",
        }
    elif label == "qConstrainedLogEI-tau1":
        return {
            "color": "brown",
            "marker": "D",
            "ls": "dashdot",
            "plot_label": "qConstrainedLogEI (tau=1.0)",
        }
    elif label == "qConstrainedLogEI-tau10":
        return {
            "color": "black",
            "marker": "s",
            "ls": "dashdot",
            "plot_label": "qConstrainedLogEI (tau=10.0)",
        }
    elif label == "qConstrainedLogEI-tau100":
        return {
            "color": "orange",
            "marker": "*",
            "ls": "dashdot",
            "plot_label": "qConstrainedLogEI (tau=100.0)",
        }
    elif label == "qConstrainedLogEI-tau0-001":
        return {
            "color": "orange",
            "marker": "o",
            "ls": "dashdot",
            "plot_label": "qConstrainedLogEI (tau=0.001)",
        }
    else:
        color = random.choice(["red", "blue", "green", "orange", "purple", "brown", "pink", "gray"])
        return {
            "color": color,
            "marker": "o",
            "ls": "solid",
            "plot_label": label,
        }


def build_plot_title(result_dir: Path) -> str:
    # return f"objective: cos(2*x)*cos(y) + sin(x), constraint: cos(x)*cos(y) - sin(x)*sin(y) - 0.5 <= 0"
    return f"objective: sin(x) + y, constraint: sin(x)*sin(y) + 0.95 <= 0"


def get_trial_count_for_plot(trial: optuna.trial.FrozenTrial) -> float:
    return trial.user_attrs.get("trial_num")


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--result-dir",
        default="./gp_simulator_results_without_evaltime_n5_bbo_dataset10",
        help="Directory containing pickled study lists.",
    )
    parser.add_argument(
        "--labels",
        nargs="+",
        required=True,
        help="Labels to load, e.g. v4.7.0 add-qlogei master",
    )
    parser.add_argument(
        "--output",
        default="async-bench-example_without_evaltime_n5_bbo_dataset10.png",
        help="Output image path.",
    )
    args = parser.parse_args()

    result_dir = Path(args.result_dir)

    _, ax = plt.subplots()

    for label in args.labels:
        study_list = load_study_list(result_dir, label)
        study_list = trim_to_same_length(study_list)

        style = get_style(label)

        plot_target_over_time(
            study_list,
            color=style["color"],
            ax=ax,
            cumtime_func=get_trial_count_for_plot,
            label=style["plot_label"],
            marker=style["marker"],
            ls=style["ls"],
            markevery=10,
        )

    ax.grid(True, which="major", alpha=0.5)
    ax.grid(True, which="minor", alpha=0.2)
    ax.set_xlabel("Number of trials")
    ax.set_ylabel("Best value so far")
    ax.set_title(build_plot_title(result_dir))
    ax.legend()

    plt.savefig(args.output, bbox_inches="tight")
    print(f"Saved to {args.output}")

if __name__ == "__main__":
    main()

Problem1

Problem2

[nabenabe0928]

Is it possible to encapsulate qLogEI instead of inheriting it?
Although inheritance works fine in Python, nested inheritance usually obfuscates codes, so I'd prefer avoid it.

[kAIto47802]

Maybe we can do so by using composition as is done in other acquisition functions for constrained optimization (i.e., ConstrainedLogEI and ConstrainedLogEHVI). This would also reduce inconsistency in the implementation strategy.

[sawa3030]

Thank you for the review! I agree that the nested inheritance was hard to follow. I've updated the implementation to encapsulate it instead.

[kAIto47802]

Suggested change

	return joint_x.sum() * 0.0 + torch.zeros(
	return torch.zeros(

This term is not needed if we multiply by 0.0

[nabenabe0928]

@sawa3030
Thank you for the change! I have a general question.
Since the following PR would make this PR shorter, do you think we should wait for the following PR to be merged first?

• Add conditional gp regressor #6721

Cc: @kAIto47802