Ax: Nonlinear parameter constraints

[yiyinying]

How to define nonlinear constraints

[lena-kashtelyan]

@yiyinying, I don't think it's currently supported, but may I learn what kind of a constraint you are looking to define, as well as whether it is a parameter constraint or an outcome metric constraint?

[yiyinying]

Should be a parameter constraint。
For example：
x1^2+x2^2=9

[Balandat]

We currently don't support general non-linear parameter constraints.

Is there a way you can reparameterize the constraint? E.g. in your example you'd have a single parameter x' in [0, 2pi] and parameterize x1 = 3 cos(x'), x2 = 3 sin(x').

May I ask what your use case is? In principle it's possible to support nonlinear parameter constraints, but it would require a significant amount of work (and it would slow down candidate generation significantly as well).

[lena-kashtelyan]

@yiyinying, does @Balandat's solution with parameterizing the constraint work for your use case? If not, could you tell us a bit more about your use case, so that we can consider supporting it?

[justinkterry]

@lena-kashtelyan @Balandat I have a practical constraint that I don't believe can be sensibly supported without nonlinear constraints: (20*n_envs*n_steps) % nminibatches == 0. That's a fairly common constraint in reinforcement learning. While work arounds exist, they're super hacky and make the search space much less smooth.

I can also explain where that constraint comes from if you'd like.

[lena-kashtelyan]

I'd be curious to hear @Balandat's thoughts on this one.

[Balandat]

So n_envs, n_steps and minibatches are all parameters here?
The modulo operator adds additional complexity since it further discretizes the search space. Assuming all of these are integers, how many values can each of them take on? If the cardinality is low enough, just enumerating the feasible points for the purpose of optimizing the acquisition function may be reasonable.

More generally, in principle non-linear constraints on the parameters are doable since generating new points already requires solving a nonlinear optimization problem. However, depending on the shape of the constraint set, the GP models we use may not work particularly well on non-convex and possibly non-connected domains.

Regarding the mechanics: We currently use SLSQP from scipy.minimize under the hood to do this. One of the challenges with non-linear constraints is that you often have to provdie the Jacobian and the constraints as a callable to the optimizer, which requires quite a bit of plumbing if we want to start from the search space. In addition, we usually transform the parameters to the unit cube for our models, and it's not trivial to automatically transform the constraints and constraint jaobians to that space.

[justinkterry]

"So n_envs, n_steps and minibatches are all parameters here?"

Yep. One is 1-4, one is 1-125, and one is 4-4096, so I can't just get a list of valid points. I turned the 4-4096 one into a 0,1 float and round the value from the model down to the nearest compatible value with the other two.

"However, depending on the shape of the constraint set, the GP models we use may not work particularly well on non-convex and possibly non-connected domains."

-Allowing for nonlinear constraints would significantly improve the convexity of a lot of problems. E.g. from a previous example of yours, if the constraint needs to be sin(x) and you accordingly transform x into sin(x) the optimization will probably be much less nice and convex. The same applies for my modulo problem.

-It wouldn't be too hard to test the shape of the space after constraint to make sure it's not pathological right?

[justinkterry]

@Balandat

[lena-kashtelyan]

We will now be tracking wishlist items / feature requests in a master issue for improved visibility: #566. Of course please feel free to still open new feature requests issues; we'll take care of thinking them through and adding them to the master issue.