Defining a Model class for ModelListGP

[Abrikosoff]

I have been trying to run a Service version of qMultiFidelityHypervolumeKnowledgeGradient; related thread I have already been able to make the input constructor work (I think!), but I keep getting stuck at the requirement for qMultiFidelityHypervolumeKnowledgeGradient to have a ModelListGP as surrogate: doing this in the GS, for example:

generation_strategy = GenerationStrategy(
                        steps=[
                            GenerationStep(
                                model=Models.SOBOL,
                                num_trials=1,  # https://github.com/facebook/Ax/issues/922
                                min_trials_observed=1,
                                max_parallelism=6,
                                model_kwargs={"seed": 9999},
                            ), 
                            GenerationStep(
                            model=Models.BOTORCH_MODULAR,
                            num_trials=-1,
                            model_kwargs={
                                "botorch_acqf_class": qMultiFidelityHypervolumeKnowledgeGradient,
                                # tried any one of these!
                                # "surrogate": Surrogate(SingleTaskGP),
                                # "surrogate": Surrogate(ModelListGP),
                                # "surrogate": ListSurrogate(SingleTaskGP)
                            },
                            model_gen_kwargs={
                                "model_gen_options": {
                                    "acqf_kwargs": {"cost_intercept": cost_intercept,
                                                    "num_fantasies": 2,
                                                    "cost_aware_utility": cost_aware_utility,
                                                    "target_fidelities": target_fidelities,
                                                    project: project,
                                                    },        
                                },
                            },
                        )
                        ]
                    )

gives me the error ValueError: qMultiFidelityHypervolumeKnowledgeGradient requires using a ModelList.

• "surrogate": Surrogate(SingleTaskGP) gives me the same error;
• "surrogate": Surrogate(ModelListGP) gives me TypeError: ModelListGP.__init__() got an unexpected keyword argument 'train_X', which is mysterious as nowhere do I have this arg in my code (but is probably because something else is trying to pass this to a surrogate object which is not correctly defined here by me).
• "surrogate": ListSurrogate(SingleTaskGP) informs me that ListSurrogate is deprecated (got this idea from this issue, which on first glance I would have thought would be the solution to this problem),

so the only thing left (I think) is via a custom model definition (like given here), ala doing something like class SimpleCustomGP(ModelListGP, GPyTorchModel):, but I am not sure how this can be properly defined (one model per MultivariateNormal?). Since I am not sure if this would be the correct way to go, I would like to ask this first before proceeding.

Thanks for help!

[bernardbeckerman]

Thanks for asking this! Let me follow up internally to see who might be best to answer.

[Balandat]

So the issue here is that by default Ax will select a (batched) non-model-list model that is incompatible with the decoupled acquisition function that qMultiFidelityHypervolumeKnowledgeGradient is a subclass of.

There is a way to avoid using batched models by passing in allow_batched_models=False as part of the surrogate specs of the model (a way of specifying model_kwargs that allow finer granular control). This would look something like passing the following (as part of model_kwargs:

"surrogate_specs": {"a": SurrogateSpec(allow_batched_models=False), "b": SurrogateSpec(allow_batched_models=False)},

However, I'm running into some weird issues with this, so I for now just hacked my way around it, see #2514 (comment) - that again surfaced a more serious limitation, see that comment.