Merge pull request #310 from nabenabe0928/add-constraints-to-var

Add constrained VaR

Author: y0z

package/samplers/value_at_risk/README.md

@@ -17,17 +17,17 @@ To implement this sampler, the author referred to [this paper](https://arxiv.org
 
 ## APIs
 
-- `RobustGPSampler(*, seed: int | None = None, independent_sampler: BaseSampler | None = None, n_startup_trials: int = 10, deterministic_objective: bool = False, constraints_func: Callable[[FrozenTrial], Sequence[float]] | None = None, warn_independent_sampling: bool = True, uniform_input_noise_ranges: dict[str, float] | None = None, normal_input_noise_stdevs: dict[str, float] | None = None)`
+- `RobustGPSampler(*, seed: int | None = None, independent_sampler: BaseSampler | None = None, n_startup_trials: int = 10, deterministic_objective: bool = False, constraints_func: Callable[[FrozenTrial], Sequence[float]] | None = None, warn_independent_sampling: bool = True, uniform_input_noise_rads: dict[str, float] | None = None, normal_input_noise_stdevs: dict[str, float] | None = None)`
   - `seed`: Random seed to initialize internal random number generator. Defaults to `None` (a seed is picked randomly).
   - `independent_sampler`: Sampler used for initial sampling (for the first `n_startup_trials` trials) and for conditional parameters. Defaults to :obj:`None` (a random sampler with the same `seed` is used).
   - `n_startup_trials`: Number of initial trials. Defaults to 10.
   - `deterministic_objective`: Whether the objective function is deterministic or not. If `True`, the sampler will fix the noise variance of the surrogate model to the minimum value (slightly above 0 to ensure numerical stability). Defaults to `False`. Currently, all the objectives will be assume to be deterministic if `True`.
   - `constraints_func`: An optional function that computes the objective constraints. It must take a `optuna.trial.FrozenTrial` and return the constraints. The return value must be a sequence of `float`. A value strictly larger than 0 means that a constraint is violated. A value equal to or smaller than 0 is considered feasible. If `constraints_func` returns more than one value for a trial, that trial is considered feasible if and only if all values are equal to 0 or smaller. The `constraints_func` will be evaluated after each successful trial. The function won't be called when trials fail or are pruned, but this behavior is subject to change in future releases.
   - `warn_independent_sampling`: If this is `True`, a warning message is emitted when the value of a parameter is sampled by using an independent sampler, meaning that no GP model is used in the sampling. Note that the parameters of the first trial in a study are always sampled via an independent sampler, so no warning messages are emitted in this case.
-  - `uniform_input_noise_ranges`: The input noise ranges for each parameter. For example, when `{"x": 0.1, "y": 0.2}`, the sampler assumes that $\\pm$ 0.1 is acceptable for `x` and $\\pm$ 0.2 is acceptable for `y`.
+  - `uniform_input_noise_rads`: The input noise radiuses for each parameter. For example, when `{"x": 0.1, "y": 0.2}`, the sampler assumes that $\\pm$ 0.1 is acceptable for `x` and $\\pm$ 0.2 is acceptable for `y`.
   - `normal_input_noise_stdevs`: The input noise standard deviations for each parameter. For example, when `{"x": 0.1, "y": 0.2}` is given, the sampler assumes that the input noise of `x` and `y` follows `N(0, 0.1**2)` and `N(0, 0.2**2)`, respectively.
 
-Please note that only one of `uniform_input_noise_ranges` and `normal_input_noise_stdevs` can be provided.
+Please note that only one of `uniform_input_noise_rads` and `normal_input_noise_stdevs` can be provided.
 
 ## Installation
 
@@ -74,7 +74,7 @@ def objective(trial: optuna.Trial) -> float:
 
 
 RobustGPSampler = optunahub.load_module("samplers/value_at_risk").RobustGPSampler
-sampler = RobustGPSampler(seed=0, uniform_input_noise_ranges={"x": 0.5, "y": 0.5})
+sampler = RobustGPSampler(seed=0, uniform_input_noise_rads={"x": 0.5, "y": 0.5})
 study = optuna.create_study(sampler=sampler)
 study.optimize(objective, n_trials=50)
 

package/samplers/value_at_risk/_gp/acqf.py

@@ -48,6 +48,39 @@ def _sample_from_normal_sobol(dim: int, n_samples: int, seed: int | None) -> tor
     return torch.erfinv(samples) * float(np.sqrt(2))
 
 
+def _sample_input_noise(
+    n_input_noise_samples: int,
+    uniform_input_noise_rads: torch.Tensor | None,
+    normal_input_noise_stdevs: torch.Tensor | None,
+    seed: int | None,
+) -> torch.Tensor:
+    assert uniform_input_noise_rads is not None or normal_input_noise_stdevs is not None
+    if normal_input_noise_stdevs is not None:
+        dim = normal_input_noise_stdevs.size(0)
+        noisy_inds = torch.where(normal_input_noise_stdevs != 0.0)
+        input_noise = torch.zeros(size=(n_input_noise_samples, dim), dtype=torch.float64)
+        input_noise[:, noisy_inds[0]] = (
+            _sample_from_normal_sobol(noisy_inds[0].size(0), n_input_noise_samples, seed)
+            * normal_input_noise_stdevs[noisy_inds]
+        )
+        return input_noise
+    elif uniform_input_noise_rads is not None:
+        dim = uniform_input_noise_rads.size(0)
+        noisy_inds = torch.where(uniform_input_noise_rads != 0.0)
+        input_noise = torch.zeros(size=(n_input_noise_samples, dim), dtype=torch.float64)
+        input_noise[:, noisy_inds[0]] = (
+            _sample_from_sobol(noisy_inds[0].size(0), n_input_noise_samples, seed)
+            * 2
+            * uniform_input_noise_rads[noisy_inds]
+            - uniform_input_noise_rads[noisy_inds]
+        )
+        return input_noise
+    else:
+        raise ValueError(
+            "Either `uniform_input_noise_rads` or `normal_input_noise_stdevs` " "must be provided."
+        )
+
+
 class BaseAcquisitionFunc(ABC):
     def __init__(self, length_scales: np.ndarray, search_space: SearchSpace) -> None:
         self.length_scales = length_scales
@@ -69,69 +102,99 @@ def eval_acqf_with_grad(self, x: np.ndarray) -> tuple[float, np.ndarray]:
         return val.item(), x_tensor.grad.detach().numpy()  # type: ignore
 
 
+class LogCumulativeProbabilityAtRisk(BaseAcquisitionFunc):
+    """The logarithm of the cumulative probability measure at risk
+
+    When we replace f(x) in VaR with 1[f(x) <= f*], the optimization of the new VaR corresponds to
+    that of the mean probability of x with input perturbation being feasible.
+    """
+
+    def __init__(
+        self,
+        gpr_list: list[GPRegressor],
+        search_space: SearchSpace,
+        confidence_level: float,
+        threshold_list: list[float],
+        n_input_noise_samples: int,
+        qmc_seed: int | None,
+        uniform_input_noise_rads: torch.Tensor | None = None,
+        normal_input_noise_stdevs: torch.Tensor | None = None,
+        stabilizing_noise: float = 1e-12,
+    ) -> None:
+        self._gpr_list = gpr_list
+        self._threshold_list = threshold_list
+        rng = np.random.RandomState(qmc_seed)
+        self._input_noise = _sample_input_noise(
+            n_input_noise_samples,
+            uniform_input_noise_rads,
+            normal_input_noise_stdevs,
+            seed=rng.random_integers(0, 2**31 - 1, size=1).item(),
+        )
+        self._stabilizing_noise = stabilizing_noise
+        self._confidence_level = confidence_level
+        super().__init__(
+            length_scales=np.mean([gpr.length_scales for gpr in gpr_list], axis=0),
+            search_space=search_space,
+        )
+
+    def eval_acqf(self, x: torch.Tensor) -> torch.Tensor:
+        x_noisy = x.unsqueeze(-2) + self._input_noise
+        log_feas_probs = torch.zeros(x_noisy.shape[:-1], dtype=torch.float64)
+        for gpr, threshold in zip(self._gpr_list, self._threshold_list):
+            means, vars_ = gpr.posterior(x_noisy)
+            sigmas = torch.sqrt(vars_ + self._stabilizing_noise)
+            # NOTE(nabenabe): integral from a to b of f(x) is integral from -b to -a of f(-x).
+            log_feas_probs += torch.special.log_ndtr((means - threshold) / sigmas)
+        n_input_noise_samples = len(self._input_noise)
+        n_risky_samples = math.ceil((1 - self._confidence_level) * n_input_noise_samples)
+        log_feas_probs_at_risk, _ = torch.topk(
+            log_feas_probs,
+            k=n_risky_samples,
+            dim=-1,
+            largest=False,
+            sorted=False,
+        )
+        return log_feas_probs_at_risk.logsumexp(dim=-1) - math.log(n_risky_samples)
+
+
 class ValueAtRisk(BaseAcquisitionFunc):
     def __init__(
         self,
         gpr: GPRegressor,
         search_space: SearchSpace,
-        alpha: float,
+        confidence_level: float,
         n_input_noise_samples: int,
         n_qmc_samples: int,
         qmc_seed: int | None,
         acqf_type: str,
-        uniform_input_noise_ranges: torch.Tensor | None = None,
+        uniform_input_noise_rads: torch.Tensor | None = None,
         normal_input_noise_stdevs: torch.Tensor | None = None,
     ) -> None:
-        assert 0 <= alpha <= 1
+        assert 0 <= confidence_level <= 1
         self._gpr = gpr
-        self._alpha = alpha
+        self._confidence_level = confidence_level
         rng = np.random.RandomState(qmc_seed)
-        self._input_noise = self._sample_input_noise(
-            n_input_noise_samples, uniform_input_noise_ranges, normal_input_noise_stdevs, rng
+        self._input_noise = _sample_input_noise(
+            n_input_noise_samples,
+            uniform_input_noise_rads,
+            normal_input_noise_stdevs,
+            seed=rng.random_integers(0, 2**31 - 1, size=1).item(),
         )
-        seed = rng.random_integers(0, 2**31 - 1, size=1).item()
         self._fixed_samples = _sample_from_normal_sobol(
-            dim=n_input_noise_samples, n_samples=n_qmc_samples, seed=seed
+            dim=n_input_noise_samples,
+            n_samples=n_qmc_samples,
+            seed=rng.random_integers(0, 2**31 - 1, size=1).item(),
         )
         self._acqf_type = acqf_type
         super().__init__(length_scales=gpr.length_scales, search_space=search_space)
 
-    @staticmethod
-    def _sample_input_noise(
-        n_input_noise_samples: int,
-        uniform_input_noise_ranges: torch.Tensor | None,
-        normal_input_noise_stdevs: torch.Tensor | None,
-        rng: np.random.RandomState,
-    ) -> torch.Tensor:
-        seed = rng.random_integers(0, 2**31 - 1, size=1).item()
-
-        def _sample_input_noise(noise_params: torch.Tensor, gen: SobolGenerator) -> torch.Tensor:
-            dim = noise_params.size(0)
-            noisy_inds = torch.where(noise_params != 0.0)
-            input_noise = torch.zeros(size=(n_input_noise_samples, dim), dtype=torch.float64)
-            input_noise[:, noisy_inds[0]] = (
-                gen(noisy_inds[0].size(0), n_input_noise_samples, seed) * noise_params[noisy_inds]
-            )
-            return input_noise
-
-        assert uniform_input_noise_ranges is not None or normal_input_noise_stdevs is not None
-        if normal_input_noise_stdevs is not None:
-            return _sample_input_noise(normal_input_noise_stdevs, _sample_from_normal_sobol)
-        elif uniform_input_noise_ranges is not None:
-            return _sample_input_noise(uniform_input_noise_ranges, _sample_from_sobol)
-        else:
-            raise ValueError(
-                "Either `uniform_input_noise_ranges` or `normal_input_noise_stdevs` "
-                "must be provided."
-            )
-
     def _value_at_risk(self, x: torch.Tensor) -> torch.Tensor:
         means, covar = self._gpr.joint_posterior(x.unsqueeze(-2) + self._input_noise)
         # TODO: Think of a better way to avoid numerical issue in the Cholesky decomposition.
         L, _ = torch.linalg.cholesky_ex(covar)
         posterior_samples = means.unsqueeze(-2) + self._fixed_samples @ L
         # If CVaR, use torch.topk instead of torch.quantile.
-        return torch.quantile(posterior_samples, q=self._alpha, dim=-1)
+        return torch.quantile(posterior_samples, q=self._confidence_level, dim=-1)
 
     def eval_acqf(self, x: torch.Tensor) -> torch.Tensor:
         """
@@ -148,3 +211,53 @@ def eval_acqf(self, x: torch.Tensor) -> torch.Tensor:
             raise NotImplementedError("NEI is not implemented yet.")
         else:
             raise ValueError(f"Unknown acqf_type: {self._acqf_type}")
+
+
+class ConstrainedLogValueAtRisk(BaseAcquisitionFunc):
+    def __init__(
+        self,
+        gpr: GPRegressor,
+        search_space: SearchSpace,
+        constraints_gpr_list: list[GPRegressor],
+        constraints_threshold_list: list[float],
+        objective_confidence_level: float,
+        feas_prob_confidence_level: float,
+        n_input_noise_samples: int,
+        n_qmc_samples: int,
+        qmc_seed: int | None,
+        acqf_type: str,
+        uniform_input_noise_rads: torch.Tensor | None = None,
+        normal_input_noise_stdevs: torch.Tensor | None = None,
+        stabilizing_noise: float = 1e-12,
+    ) -> None:
+        self._value_at_risk = ValueAtRisk(
+            gpr=gpr,
+            search_space=search_space,
+            confidence_level=objective_confidence_level,
+            n_input_noise_samples=n_input_noise_samples,
+            n_qmc_samples=n_qmc_samples,
+            qmc_seed=qmc_seed,
+            acqf_type=acqf_type,
+            uniform_input_noise_rads=uniform_input_noise_rads,
+            normal_input_noise_stdevs=normal_input_noise_stdevs,
+        )
+        self._log_prob_at_risk = LogCumulativeProbabilityAtRisk(
+            gpr_list=constraints_gpr_list,
+            search_space=search_space,
+            confidence_level=feas_prob_confidence_level,
+            threshold_list=constraints_threshold_list,
+            n_input_noise_samples=n_input_noise_samples,
+            qmc_seed=qmc_seed,
+            uniform_input_noise_rads=uniform_input_noise_rads,
+            normal_input_noise_stdevs=normal_input_noise_stdevs,
+            stabilizing_noise=stabilizing_noise,
+        )
+        assert torch.allclose(
+            self._log_prob_at_risk._input_noise, self._value_at_risk._input_noise
+        )
+        super().__init__(self._value_at_risk.length_scales, search_space=search_space)
+
+    def eval_acqf(self, x: torch.Tensor) -> torch.Tensor:
+        return self._value_at_risk.eval_acqf(x).clamp_min_(
+            _EPS
+        ).log_() + self._log_prob_at_risk.eval_acqf(x)

package/samplers/value_at_risk/example.py

@@ -25,6 +25,6 @@ def objective(trial: optuna.Trial) -> float:
 
 
 RobustGPSampler = optunahub.load_module("samplers/value_at_risk").RobustGPSampler
-sampler = RobustGPSampler(seed=0, uniform_input_noise_ranges={"x": 0.5, "y": 0.5})
+sampler = RobustGPSampler(seed=0, uniform_input_noise_rads={"x": 0.5, "y": 0.5})
 study = optuna.create_study(sampler=sampler)
 study.optimize(objective, n_trials=50)