[GENERAL SUPPORT]: Hooking up custom BoTorch acquisition functions

[parrangoiz]

Question

Hello, general question here. I'm working on hooking up the ECI acquisition function from the BoTorch tutorials to the Ax machinery. ECI is not part of the official BoTorch library as far as I can tell and only exists as an example/tutorial, but even that demo implementation there seems to do a good job at solving the problem I am working on. So now I want to hook it up to Ax for easier data management, saving experiments, etc. I've gone through all the required steps (mainly defining an input constructor) and have run into a couple speed bumps along the way:

1. The Ax Acquisition wrapper class attempts to inject many arguments (see here) which may or may not be input arguments to the custom BoTorch acquisition class we're trying to plug into. For instance, in this case a posterior_transform argument is passed, so I ended up having to add it as well to the ECI implementation even though it is not used for anything. If I don't add it, the decorated input constructor, which gets preprocessed by allow_only_specific_variable_kwargs, throws an error.
2. The Acquisition constructor preprocesses the outcome constraints in the TorchOptConfig container into a list of Callables. It seems that many BoTorch acquisition functions (e.g. most of the ones in botorch.acquisition.monte_carlo) take in outcome constraints in this form, but others don't (e.g. some of the analytic functions that support constraints). The toy implementation of ECI assumed a different constraint format, but in this case it's simple enough for me to modify the implementation to work with callables instead.

In sum, so far it looks like I am managing to establish the connection, albeit in a somewhat hacky way. Thoughts? Am I on the right track? More generally I get the feeling that Ax's Acquisition already makes too many assumptions and therefore subclassing it to specialize it to custom backend implementations is not a good approach, and one has to resort to hacking a bit. I saw some discussion in past issues regarding this problem and the need to refactor the API so maybe that's a good long-ish term project to increase flexibility. For now, I am happy to (eventually) put in a BoTorch MR to make ECI part of the official library if there is interest in this and once I get some pointers/feedback on my approach.

Please provide any relevant code snippet if applicable.

Code of Conduct

• I agree to follow this Ax's Code of Conduct

[Balandat]

@parrangoiz thanks for sharing your feedback on the API - at a high level there is a pretty strong tension between having a setup that is easily modifiable and that doesn't make many assumptions in order to enable fast research iteration (this is BoTorch's design philosophy) on one side, and a setup that is strict and explicit to avoid bugs and simplify things such as serialization/storage (this is Ax's design philosophy) on the other.

The Ax Acquisition wrapper class attempts to inject many arguments (see here) which may or may not be input arguments to the custom BoTorch acquisition class we're trying to plug into. For instance, in this case a posterior_transform argument is passed, so I ended up having to add it as well to the ECI implementation even though it is not used for anything. If I don't add it, the decorated input constructor, which gets preprocessed by allow_only_specific_variable_kwargs, throws an error.

The reason we're strict about the arguments used in input constructors (and others) is that in the past - when we were not - we ended up introducing a number of pretty nasty bugs that stemmed from silently ignoring such kwargs. I think it does make sense to be strict here, but we can probably do a better job at explaining this / adding documentation to make this less of a gotcha.

The Acquisition constructor preprocesses the outcome constraints in the TorchOptConfig container into a list of Callables. It seems that many BoTorch acquisition functions (e.g. most of the ones in botorch.acquisition.monte_carlo) take in outcome constraints in this form, but others don't (e.g. some of the analytic functions that support constraints). The toy implementation of ECI assumed a different constraint format, but in this case it's simple enough for me to modify the implementation to work with callables instead.

The tradeoff here is between how much work do we want to do at a higher level vs. at lower levels. Having a generic high-level representation of the constraints (e.g. not callables) would provide more flexibility but also would require a lot more boilerplate code for most if not all of the existing acquisition functions. There are a couple of ways to deal with that - one would be to use a more flexible format for the API, package the conversion into some helper function, and then just call that in each input constructor. Alternatively one could have two separate (say Acquisition and DefaultBoTorchAcquisition) classes that deal with the constraints differently. The latter is somewhat dangerous though b/c it likely will lead to many such classes popping up and making things harder to understand.

In sum, so far it looks like I am managing to establish the connection, albeit in a somewhat hacky way. Thoughts? Am I on the right track?

You're definitely on the right track, this is how one does hook up custom acquisition functions.

More generally I get the feeling that Ax's Acquisition already makes too many assumptions and therefore subclassing it to specialize it to custom backend implementations is not a good approach, and one has to resort to hacking a bit. I saw some discussion in past issues regarding this problem and the need to refactor the API so maybe that's a good long-ish term project to increase flexibility.

Indeed, this could potentially be a good investment. We have some more pressing short term work but in the mid-to-long term this seems possible. It's great to see community engagement like yours to hook things up and we definitely want to encourage that and make it easier to contribute.

For now, I am happy to (eventually) put in a BoTorch MR to make ECI part of the official library if there is interest in this and once I get some pointers/feedback on my approach.

Yes, we'd love to incorporate ECI into BoTorch (+Ax) and we're happy to provide feedback on the design and implementation!