[Feature Request] Add a generation strategy config to ax-sweeper plugin

[lwelzel]

🚀 Feature Request

Add the possibility to configure the generation strategy of the sweeper in the ax-sweeper plugin.

Motivation

I want to specify a specific generation strategy for the sweep executed by the AxSweeper. This is important when I want to use the sweeper to generate samples in a specific way, e.g. only on a Sobol grid. Defining a specific generation strategy is necessary when the default is not good, e.g. when running large sweeps that do not need to optimize an objective, or when the automatically chosen strategy is not good.
Right now the generation strategy is automatically chosen by the AxClient. From what I understand it is not possible to change this using a hydra config file.

This is also related to this feature request for setting max_parallelism (issue #1837) since max_parallelism is defined in the generation steps of a generation strategy. The ax-sweeper is also now several versions behind ax-platform and botorch so that it would be nice to also update the overall ax-sweeper interface (#2777).

Pitch

I want to be able to define a generation strategy in a config file, e.g. like:

defaults:
  - override hydra/sweeper: ax
hydra:
  sweep:
    dir: ${hydra.job.name}_${now:%Y-%m-%d_%H-%M-%S}
    subdir: jobnr_${hydra.job.num}_jobid_${hydra.job.id}
  sweeper:
    ax_config:
      experiment:
        name: "example"
      generation_strategy: # new generation_strategy config as part of ax_config
        steps: # can define several steps on the generation strategy
                   # I should be able to set the kwargs of ax.modelbridge.generation_strategy.GenerationStep in this config
        - model: SOBOL # new (list of) generation_step config as part of generation_strategy config
          num_trials: 5
          max_parallelism: 1  # the max_parallelism for a step can be set
        - model: GPEI
          num_trials: 2
        name: default_strategy
      max_trials: 7 # this should automatically be set to the sum of the num_trials of the steps
      params:
        ...

Describe alternatives you've considered
The other sweepers don't support the same features as Ax/BOTorch so that I would need to write something that replaces the ax-sweeper plugin.

Are you willing to open a pull request?
I have implemented a possible solution (PR #2978), however, I have not added tests for the feature yet since I dont have experience with defining automated tests.