[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!

[parrangoiz]

Thanks for the detailed response!

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.

In this specific scenario I am now realizing that there's a method I can specialize that would take care of the posterior_transform argument (for instance just making it None if I don't need it). Nevertheless the possibility of more footguns is still there. I agree this could be mitigated to some extent by more clearly explaining in the docs (a) which arguments will be injected by Ax's Acquisition to the BoTorch-side acquisition function (and therefore should be placed there by the user), and (b) which exceptions are allowed, i.e. the ones listed here. I'm happy to give it a stab updating the docs (maybe here?) but the main challenge I see is that the set of arguments that are constructed and passed is not static but is decided by some pre-programmed logic, so at some point it becomes easier for the user to just read the source code to figure out exactly what arguments their custom class should take. I'll let the issue rest but maybe this discussion could be helpful during a future API refactor : )

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.

This could be reasonable and could help create a more uniform API. My main concern is that some BoTorch acquisition functions take in constraints as a list of callables whereas others take them in the form {i: [lower, upper], ...} so if I'm understanding the code correctly only the former currently work with Ax out of the box.

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

Sweet! I can't promise I'll be able to get to this super soon but I'm very interested in getting this merged at some point because as I said ECI is already looking quite useful for my research. Just for future reference, how do you normally go about designing/implementing large-ish features like this? Would I just put in a PR with a rough draft and go from there?

[Balandat]

Just for future reference, how do you normally go about designing/implementing large-ish features like this? Would I just put in a PR with a rough draft and go from there?

We usually start off with a relatively high level design - this can just be a written summary and some pseudocode; possibly on an issue opened for that feature request in particular. Then the team can provide some initial feedback on the direction and high-level design, iterate (if necessary), and the next step then would be a draft PR. That way we can provide early feedback without having people invest a lot of time into coding something up in a suboptimal way.

[parrangoiz]

Makes sense thanks. I'll put in an issue when I get a chance.