Rounding issues when generating with constraints

[yoavnavon]

Problem

Hi, I've been trying to generate trials for an experiment with constraints, and I keep getting an error because the proposed trials don't satisfy the constraints by small differences, probably due to some rounding error somewhere. I did submit a similar issue once (#439) that seemed to be resolved, but may be related.

How To Replicate Error

from ax.service.ax_client import AxClient
from ax.modelbridge import get_sobol
from ax.modelbridge.registry import Models

client = AxClient()
client.create_experiment(
    name="experiment",
    parameters=[
        {"name": "x", "type": "range", "bounds": [0, 10], "value_type": "float"},
        {"name": "y", "type": "range", "bounds": [0, 10], "value_type": "float"},
    ],
    parameter_constraints=[
        "x + y <= 15",
        "x + y >= 10"
        ],
)

# Generate and attach sobol points
search_space = client.experiment.search_space
sobol = get_sobol(search_space, deduplicate=True, seed=0)
sobol_points = sobol.gen(5)
for i, arm in enumerate(sobol_points.arms):
    params = arm.parameters
    _, idx = client.attach_trial(params)
    client.complete_trial(idx, params["x"])

# Generate and attach from GP
bopt = Models.BOTORCH(
    experiment=client.experiment,
    data=client.experiment.fetch_data(),
)
samples = bopt.gen(5).param_df.reset_index(drop=True).to_dict("records")
for sample in samples:
    _, idx = client.attach_trial(sample)
    client.complete_trial(idx, 0.5)

Stack Trace

[INFO 11-29 15:11:39] ax.service.ax_client: Starting optimization with verbose logging. To disable logging, set the `verbose_logging` argument to `False`. Note that float values in the logs are rounded to 2 decimal points.
[INFO 11-29 15:11:39] ax.modelbridge.dispatch_utils: Using Bayesian optimization since there are more ordered parameters than there are categories for the unordered categorical parameters.
[INFO 11-29 15:11:39] ax.modelbridge.dispatch_utils: Using Bayesian Optimization generation strategy: GenerationStrategy(name='Sobol+GPEI', steps=[Sobol for 5 trials, GPEI for subsequent trials]). Iterations after 5 will take longer to generate due to  model-fitting.
[INFO 11-29 15:11:39] ax.service.ax_client: Attached custom parameterization {'x': 4.75, 'y': 5.93} as trial 0.
[INFO 11-29 15:11:39] ax.service.ax_client: Completed trial 0 with data: {'objective': (4.75, None)}.
[INFO 11-29 15:11:39] ax.service.ax_client: Attached custom parameterization {'x': 2.15, 'y': 9.74} as trial 1.
[INFO 11-29 15:11:39] ax.service.ax_client: Completed trial 1 with data: {'objective': (2.15, None)}.
[INFO 11-29 15:11:39] ax.service.ax_client: Attached custom parameterization {'x': 8.54, 'y': 4.03} as trial 2.
[INFO 11-29 15:11:39] ax.service.ax_client: Completed trial 2 with data: {'objective': (8.54, None)}.
[INFO 11-29 15:11:39] ax.service.ax_client: Attached custom parameterization {'x': 7.42, 'y': 6.95} as trial 3.
[INFO 11-29 15:11:39] ax.service.ax_client: Completed trial 3 with data: {'objective': (7.42, None)}.
[INFO 11-29 15:11:39] ax.service.ax_client: Attached custom parameterization {'x': 2.58, 'y': 7.87} as trial 4.
[INFO 11-29 15:11:39] ax.service.ax_client: Completed trial 4 with data: {'objective': (2.58, None)}.
{'x': 9.999999999999755, 'y': 5.000000000001556}
Traceback (most recent call last):
  File "error.py", line 35, in <module>
    _, idx = client.attach_trial(sample)
  File "/Users/yoavnavon/Documents/NotCo/Bridge/toolbox-backend-2/.venv/lib/python3.8/site-packages/ax/service/ax_client.py", line 620, in attach_trial
    self._validate_search_space_membership(parameters=parameters)
  File "/Users/yoavnavon/Documents/NotCo/Bridge/toolbox-backend-2/.venv/lib/python3.8/site-packages/ax/service/ax_client.py", line 1375, in _validate_search_space_membership
    self.experiment.search_space.check_membership(
  File "/Users/yoavnavon/Documents/NotCo/Bridge/toolbox-backend-2/.venv/lib/python3.8/site-packages/ax/core/search_space.py", line 181, in check_membership
    raise ValueError(f"Parameter constraint {constraint} is violated.")
ValueError: Parameter constraint ParameterConstraint(1.0*x + 1.0*y <= 15.0) is violated.

[lena-kashtelyan]

Hi @yoavnavon, we'll look into this, but before we do, I have a quick question. What is the reason why you are manually generating and attaching trials instead of using ax_client.get_next_trial and ax_client.complete_trial? See this tutorial for full example: https://ax.dev/tutorials/gpei_hartmann_service.html.

I think using the canonical AxClient flow should resolve this issue for you.

[yoavnavon]

Hi @lena-kashtelyan, so I simplified the example the most I could, but basically I needed a lot of customization on the GP and acquisition function level that I could not do with the GenerationStrategy & GenerationStep. I needed to pass a bunch of arguments like this:

bopt = Models.BOTORCH(
    experiment=client.experiment,
    data=client.experiment.fetch_data(),
    acqf_constructor=acqf_func,
    model_constructor=get_and_fit_OneHotGP,
    kwargs=gp_kwargs,
)
samples = model.gen(
    n,
    model_gen_options={
        "pending_observations": pending,
        "acquisition_function_kwargs": {"beta": beta},
    },
    fixed_features=fixed_features,
)

I realize now that at the time I was using a somewhat old version of Ax (0.1.14), so maybe those problems were already addressed, but it would be helpful to know why it fails with my way.

[lena-kashtelyan]

If you are not looking to use Ax storage (JSON or SQL), you can actually pass all those arguments through GenerationStep and GenerationStrategy, with exception of perhaps the fixed features (but do you need to be passing those manually? why?)

To achieve your code snippet above, you would do:

gs = GenerationStrategy(
    steps=[
        ...,  # Sobol initialization step
        GenerationStep(
             model=Models.BOTORCH,
             model_kwargs={
                 "acqf_constructor": acqf_func,
                 "model_constructor": get_and_fit_OneHotGP,
                 ...  # Any other kwargs you need passed to `Models.BOTORCH`; data to model will be passed automatically
            },
             model_gen_kwargs={
                 "acquisition_function_kwargs": {"beta": beta},
                  # If your fixed features are predetermined, you could pass them here (or just let Ax determine them?)
                 "fixed_features": ... 
             },    
        )
    ]
)

ax_client = AxClient(generation_strategy=gs)

For more detail on generation strategy, we have this relatively new tutorial: https://ax.dev/tutorials/generation_strategy.html : )

Let me know if this doesn't work for you, we can then try to look into the rounding issue, but generally attaching trials that violate constraints is not supported through AxClient, so I'm not sure if we can easily fix that. It would be best to use the Service API canonically instead.

---

P.S. The reason why this will prohibit use of Ax storage is that some of the arguments you are specifying via model_kwargs will not be easily serializable. However, if you will need to be using storage with this setup, let us know, there is still a way, just not the usual and simplest one : )

P.S.2 Also we have this new modular setup for making BoTorch models in Ax, see this tutorial: https://ax.dev/tutorials/modular_botax.html. This should allow you to achieve the same thing you are currently achieving with passing acqf_func and get_and_fit_OneHotGP, but in a way that will be fully supported by Ax storage and will let your generation strategy be resumable from JSON or SQL.

[Balandat]

It looks like here the constraint not being satisfied is just an issue of numerical inaccuracy. I wonder if we should just change

Ax/ax/core/parameter_constraint.py

Line 78 in 4929a2e

return weighted_sum <= self._bound

to sth like weighted_sum <= self._bound + eps where ideally eps is dynamically looked up as the machine epsilon (e.g. via numpy). I can't really see much downsides to be just a little bit forgiving here

[Balandat]

The reason #439 did not address this was b/c this was purely on the bounds on an individual parameter, here the issue is with the bound on a linear constraint.

[Balandat]

#734 should deal with this.

[lena-kashtelyan]

Fix for this is now on master and will be part of the next stable release!

[yoavnavon]

Thanks for your help!

[lena-kashtelyan]

The fix is now part of the latest stable version: 0.2.3.