Change how noisy constant values are handled

- Allow specifying nominal values of noisy constants explicitly
  instead of using the center value of the interval.

- Sample noisy constant values first and then optimize other parameter's
  values based on those values, instead of optimizing optimizing other
  parameters assuming the noisy constant's nominal values, and then
  changing the noisy constant's value afterwords.
  I believe this is desirable because the sampler can explore desirable
  parameter values under different noisy constant values.

Author: msakai

package/samplers/value_at_risk/_gp/acqf.py

@@ -118,6 +118,8 @@ def __init__(
         threshold_list: list[float],
         n_input_noise_samples: int,
         qmc_seed: int | None,
+        fixed_indices: torch.Tensor,
+        fixed_values: torch.Tensor,
         uniform_input_noise_rads: torch.Tensor | None = None,
         normal_input_noise_stdevs: torch.Tensor | None = None,
         stabilizing_noise: float = 1e-12,
@@ -133,12 +135,22 @@ def __init__(
         )
         self._stabilizing_noise = stabilizing_noise
         self._confidence_level = confidence_level
+        self._fixed_indices = fixed_indices
+        self._fixed_values = fixed_values
         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:
+        # The search space of constant noisy parameters is internally replaced with IntDistribution(0, 0),
+        # so that their normalized values passed as x is always 0.5. However, values passed to
+        # const_noisy_param_values argument of RobustGPSampler._get_value_at_risk may be normalized to
+        # different values under the original search space used by GPRegressor.  So we carry around the
+        # normalized version of const_noisy_param_values explicity and use them instead.
+        x = x.clone()
+        x[:, self._fixed_indices] = self._fixed_values
+
         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):
@@ -168,6 +180,8 @@ def __init__(
         n_qmc_samples: int,
         qmc_seed: int | None,
         acqf_type: str,
+        fixed_indices: torch.Tensor,
+        fixed_values: torch.Tensor,
         uniform_input_noise_rads: torch.Tensor | None = None,
         normal_input_noise_stdevs: torch.Tensor | None = None,
     ) -> None:
@@ -187,6 +201,8 @@ def __init__(
             seed=rng.random_integers(0, 2**31 - 1, size=1).item(),
         )
         self._acqf_type = acqf_type
+        self._fixed_indices = fixed_indices
+        self._fixed_values = fixed_values
         super().__init__(length_scales=gpr.length_scales, search_space=search_space)
 
     def _value_at_risk(self, x: torch.Tensor) -> torch.Tensor:
@@ -206,6 +222,15 @@ def eval_acqf(self, x: torch.Tensor) -> torch.Tensor:
         4. Then compute (mc_value_at_risk - f0).clamp_min(0).mean()
         Appendix B.2 of https://www.robots.ox.ac.uk/~mosb/public/pdf/136/full_thesis.pdf
         """
+
+        # The search space of constant noisy parameters is internally replaced with IntDistribution(0, 0),
+        # so that their normalized values passed as x is always 0.5. However, values passed to
+        # const_noisy_param_values argument of RobustGPSampler._get_value_at_risk may be normalized to
+        # different values under the original search space used by GPRegressor.  So we carry around the
+        # normalized version of const_noisy_param_values explicity and use them instead.
+        x = x.clone()
+        x[:, self._fixed_indices] = self._fixed_values
+
         if self._acqf_type == "mean":
             return self._value_at_risk(x).mean(dim=-1)
         elif self._acqf_type == "nei":
@@ -227,6 +252,8 @@ def __init__(
         n_qmc_samples: int,
         qmc_seed: int | None,
         acqf_type: str,
+        fixed_indices: torch.Tensor,
+        fixed_values: torch.Tensor,
         uniform_input_noise_rads: torch.Tensor | None = None,
         normal_input_noise_stdevs: torch.Tensor | None = None,
         stabilizing_noise: float = 1e-12,
@@ -241,6 +268,8 @@ def __init__(
             acqf_type=acqf_type,
             uniform_input_noise_rads=uniform_input_noise_rads,
             normal_input_noise_stdevs=normal_input_noise_stdevs,
+            fixed_indices=fixed_indices,
+            fixed_values=fixed_values,
         )
         self._log_prob_at_risk = LogCumulativeProbabilityAtRisk(
             gpr_list=constraints_gpr_list,
@@ -252,6 +281,8 @@ def __init__(
             uniform_input_noise_rads=uniform_input_noise_rads,
             normal_input_noise_stdevs=normal_input_noise_stdevs,
             stabilizing_noise=stabilizing_noise,
+            fixed_indices=fixed_indices,
+            fixed_values=fixed_values,
         )
         assert torch.allclose(
             self._log_prob_at_risk._input_noise, self._value_at_risk._input_noise

package/samplers/value_at_risk/sampler.py

@@ -120,6 +120,7 @@ def __init__(
         uniform_input_noise_rads: dict[str, float] | None = None,
         normal_input_noise_stdevs: dict[str, float] | None = None,
         const_noisy_param_names: list[str] | None = None,
+        const_noisy_param_nominal_values: dict[str, float] | None = None,
     ) -> None:
         if uniform_input_noise_rads is None and normal_input_noise_stdevs is None:
             raise ValueError(
@@ -150,6 +151,7 @@ def __init__(
         self._uniform_input_noise_rads = uniform_input_noise_rads
         self._normal_input_noise_stdevs = normal_input_noise_stdevs
         self._const_noisy_param_names = const_noisy_param_names or []
+        self._const_noisy_param_nominal_values = const_noisy_param_nominal_values or {}
         self._rng = LazyRandomState(seed)
         self._independent_sampler = independent_sampler or optuna.samplers.RandomSampler(seed=seed)
         self._intersection_search_space = optuna.search_space.IntersectionSearchSpace()
@@ -262,7 +264,8 @@ def _get_constraints_acqf_args(
         return constraints_gprs, constraints_threshold_list
 
     def _get_internal_search_space_with_fixed_params(
-        self, search_space: dict[str, BaseDistribution]
+        self,
+        search_space: dict[str, BaseDistribution],
     ) -> gp_search_space.SearchSpace:
         search_space_with_fixed_params = search_space.copy()
         for param_name in self._const_noisy_param_names:
@@ -275,6 +278,7 @@ def _get_value_at_risk(
         internal_search_space: gp_search_space.SearchSpace,
         search_space: dict[str, BaseDistribution],
         acqf_type: str,
+        const_noisy_param_values: dict[str, float],
         constraints_gpr_list: list[gp.GPRegressor] | None = None,
         constraints_threshold_list: list[float] | None = None,
     ) -> acqf_module.ValueAtRisk | acqf_module.ConstrainedLogValueAtRisk:
@@ -313,17 +317,37 @@ def _get_scaled_input_noise_params(
             for i, param_name in enumerate(search_space)
             if param_name in self._const_noisy_param_names
         ]
+
+        def normalize(dist: BaseDistribution, x: float) -> float:
+            assert isinstance(
+                dist,
+                (optuna.distributions.IntDistribution, optuna.distributions.FloatDistribution),
+            )
+            return (x - dist.low) / (dist.high - dist.low)
+
+        const_noisy_param_normalized_values = [
+            normalize(dist, const_noisy_param_values[param_name])
+            if param_name in const_noisy_param_values
+            else 0.5
+            for i, (param_name, dist) in enumerate(search_space.items())
+            if param_name in self._const_noisy_param_names
+        ]
+
         if self._uniform_input_noise_rads is not None:
             scaled_input_noise_params = _get_scaled_input_noise_params(
                 self._uniform_input_noise_rads, "uniform_input_noise_rads"
             )
+            # FIXME(sakai): If the fixed value is not at the center of the range,
+            # \pm 0.5 may not cover the domain.
             scaled_input_noise_params[const_noise_param_inds] = 0.5
             noise_kwargs["uniform_input_noise_rads"] = scaled_input_noise_params
         elif self._normal_input_noise_stdevs is not None:
             scaled_input_noise_params = _get_scaled_input_noise_params(
                 self._normal_input_noise_stdevs, "normal_input_noise_stdevs"
             )
             # NOTE(nabenabe): \pm 2 sigma will cover the domain.
+            # FIXME(sakai): If the fixed value is not at the center of the range,
+            # \pm 2 sigma may not cover the domain.
             scaled_input_noise_params[const_noise_param_inds] = 0.25
             noise_kwargs["normal_input_noise_stdevs"] = scaled_input_noise_params
         else:
@@ -341,6 +365,10 @@ def _get_scaled_input_noise_params(
                 n_qmc_samples=self._n_qmc_samples,
                 qmc_seed=self._rng.rng.randint(1 << 30),
                 acqf_type=acqf_type,
+                fixed_indices=torch.tensor(const_noise_param_inds, dtype=torch.int64),
+                fixed_values=torch.tensor(
+                    const_noisy_param_normalized_values, dtype=torch.float64
+                ),
                 **noise_kwargs,
             )
         else:
@@ -355,6 +383,10 @@ def _get_scaled_input_noise_params(
                 n_qmc_samples=self._n_qmc_samples,
                 qmc_seed=self._rng.rng.randint(1 << 30),
                 acqf_type=acqf_type,
+                fixed_indices=torch.tensor(const_noise_param_inds, dtype=torch.int64),
+                fixed_values=torch.tensor(
+                    const_noisy_param_normalized_values, dtype=torch.float64
+                ),
                 **noise_kwargs,
             )
 
@@ -419,7 +451,11 @@ def _get_gpr_list(
         return gprs_list
 
     def _optimize_params(
-        self, study: Study, trials: list[FrozenTrial], search_space: dict[str, BaseDistribution]
+        self,
+        study: Study,
+        trials: list[FrozenTrial],
+        search_space: dict[str, BaseDistribution],
+        const_noisy_param_values: dict[str, float],
     ) -> dict[str, Any]:
         if search_space == {}:
             return {}
@@ -437,6 +473,7 @@ def _optimize_params(
                 internal_search_space,
                 search_space,
                 acqf_type="mean",
+                const_noisy_param_values=const_noisy_param_values,
             )
         else:
             constraint_vals, _ = _get_constraint_vals_and_feasibility(study, trials)
@@ -453,10 +490,18 @@ def _optimize_params(
                 acqf_type="mean",
                 constraints_gpr_list=constr_gpr_list,
                 constraints_threshold_list=constr_threshold_list,
+                const_noisy_param_values=const_noisy_param_values,
             )
 
         normalized_param = self._optimize_acqf(acqf)
-        return internal_search_space.get_unnormalized_param(normalized_param)
+        # The normalized values of constant noise parameters are fixed at 0.5 during search
+        # regardless of their original values given as const_noisy_param_values, so
+        # `internal_search_space.get_unnormalized_param` cannot decode them correctly.
+        # Therefore, we overwrite those values with their original values.
+        return (
+            internal_search_space.get_unnormalized_param(normalized_param)
+            | const_noisy_param_values
+        )
 
     def sample_relative(
         self, study: Study, trial: FrozenTrial, search_space: dict[str, BaseDistribution]
@@ -465,15 +510,14 @@ def sample_relative(
         if len(trials) < self._n_startup_trials:
             return {}
 
-        params = self._optimize_params(study, trials, search_space)
-
         # Perturb constant noisy parameter uniformly
+        const_noisy_param_values = {}
         for name in self._const_noisy_param_names:
             dist = search_space[name]
             assert isinstance(dist, optuna.distributions.FloatDistribution)
-            params[name] = self._rng.rng.uniform(dist.low, dist.high)
+            const_noisy_param_values[name] = self._rng.rng.uniform(dist.low, dist.high)
 
-        return params
+        return self._optimize_params(study, trials, search_space, const_noisy_param_values)
 
     def get_robust_trial(self, study: Study) -> FrozenTrial:
         states = (TrialState.COMPLETE,)
@@ -485,7 +529,11 @@ def get_robust_trial(self, study: Study) -> FrozenTrial:
         acqf: acqf_module.BaseAcquisitionFunc
         if self._constraints_func is None:
             acqf = self._get_value_at_risk(
-                gpr, internal_search_space, search_space, acqf_type="mean"
+                gpr,
+                internal_search_space,
+                search_space,
+                acqf_type="mean",
+                const_noisy_param_values=self._const_noisy_param_nominal_values,
             )
         else:
             constraint_vals, _ = _get_constraint_vals_and_feasibility(study, trials)
@@ -502,6 +550,7 @@ def get_robust_trial(self, study: Study) -> FrozenTrial:
                 acqf_type="mean",
                 constraints_gpr_list=constr_gpr_list,
                 constraints_threshold_list=constr_threshold_list,
+                const_noisy_param_values=self._const_noisy_param_nominal_values,
             )
 
         best_idx = np.argmax(acqf.eval_acqf_no_grad(X_train)).item()
@@ -511,7 +560,9 @@ def get_robust_params(self, study: Study) -> dict[str, Any]:
         states = (TrialState.COMPLETE,)
         trials = study._get_trials(deepcopy=False, states=states, use_cache=True)
         search_space = self.infer_relative_search_space(study, trials[0])
-        return self._optimize_params(study, trials, search_space)
+        return self._optimize_params(
+            study, trials, search_space, self._const_noisy_param_nominal_values
+        )
 
     def sample_independent(
         self,