Address Issue #84 .

kge/config-default.yaml

@@ -460,6 +460,10 @@ ax_search:
  # If equal or larger than num_trials, only Sobal trials will be run.
  num_sobol_trials: -1
 
+ # Random seed for generating the sobol sequence. Has to be fixed for each
+ # experiment, or else resuming the sobol sequence is inconsistent.
+ sobol_seed: 0
+
  # Search space definition passed to ax. See create_experiment in
  # https://ax.dev/api/service.html#module-ax.service.ax_client
  parameters: []

kge/job/ax_search.py

@@ -40,7 +40,7 @@ def init_search(self):
  num_trials=self.num_sobol_trials,
  min_trials_observed=ceil(self.num_sobol_trials / 2),
  enforce_num_trials=True,
- model_kwargs={'seed': 0}
+ model_kwargs={"seed": self.config.get("ax_search.sobol_seed")},
  ),
  GenerationStep(
  model=Models.GPEI,
@@ -62,15 +62,21 @@ def init_search(self):
  # END: from /ax/service/utils/dispatch.py
 
  self.ax_client = AxClient(generation_strategy=generation_strategy)
+ choose_generation_strategy_kwargs=dict()
  else:
  self.ax_client = AxClient()
+ # set random_seed that will be used by auto created sobol search from ax
+ # note that here the argument is called "random_seed" not "seed"
+ choose_generation_strategy_kwargs={
+ "random_seed": self.config.get("ax_search.sobol_seed")
+ },
  self.ax_client.create_experiment(
  name=self.job_id,
  parameters=self.config.get("ax_search.parameters"),
  objective_name="metric_value",
  minimize=False,
  parameter_constraints=self.config.get("ax_search.parameter_constraints"),
- choose_generation_strategy_kwargs={'random_seed': 0},
+ choose_generation_strategy_kwargs=choose_generation_strategy_kwargs,
  )
  self.config.log(
  "ax search initialized with {}".format(self.ax_client.generation_strategy)