[Feature Request] Specifying fixed features to AxClient.get_next_trial (to conduct contextual BO using Service API)

[MattiaVerticchio]

Hi everyone. First of all, thanks for working on Ax! 😁

I have a setting where I’d like to optimize a function based on a certain context. The context is, in this case, the day of the week.

If I understood correctly, the only way to have contextual BO in Ax is by using the Developer API; setting the context as input covariates, and then generating a new_trial with fixed_features:

gpei = Models.BOTORCH(
    experiment=exp,
    data=exp.fetch_data()
)

generator_run = gpei.gen(n=1)

fixed_features = ObservationFeatures(parameters={"day_of_week": 0})

trial = exp.new_trial(
    generator_run=generator_run,
    fixed_features=fixed_features)
)

For this kind of task, the Developer API is more flexible but adds a lot of complexity. Is there any other way to do contextual BO, or pass fixed_features when using the simpler Service API?

Thanks!

[lena-kashtelyan]

Hi @MattiaVerticchio, I think right now the Service API does not support this case, but it definitely could –– it seems reasonable for get_next_trial to be able to accept fixed features. I'll mark this as a feature request!

In the meantime, you can hack around this like in this comment: #733 (comment), so basically generating the parameterizations outside of AxClient via Models.BOTORCH and then attaching them via attach_trial. Generally we really don't recommend this approach, but while the functionality you'd like isn't supported, you can make use of this perhaps?

[MattiaVerticchio]

In the meantime, you can hack around this like in this comment: #733 (comment), so basically generating the parameterizations outside of AxClient via Models.BOTORCH and then attaching them via attach_trial. Generally we really don't recommend this approach, but while the functionality you'd like isn't supported, you can make use of this perhaps?

Yes, that’s definitely better, thanks for your help!

[lena-kashtelyan]

Putting this on wishlist for now, but we will likely support this soon!

[sgbaird]

@MattiaVerticchio, do you mind including a reproducer for this example if you have one available?