Issues with `optimize`: needs more documentation, maybe does not return the correct best parameters?

[kurtosis]

I'm running the Get Started example with optimize and seeing a couple issues:

1. Setting random_seed does not lead to reproducible output. Values only change by a small amount on subsequent runs so it's not 100% "random"- perhaps there's another RNG that feeds into one part of this?
2. best_parameters do not actually come from the trial with the best result. In the example below, they come from trial 14, while the actual best result was in trial 11. (This issue persisted even though the output values varied when I reran it.)

In [13]: %paste                                                                                                                               
import ax
print(ax.__version__)
best_parameters, best_values, experiment, model = ax.optimize(
        parameters=[
          {
            "name": "x1",
            "type": "range",
            "bounds": [-10.0, 10.0],
          },
          {
            "name": "x2",
            "type": "range",
            "bounds": [-10.0, 10.0],
          },
        ],
        # Booth function
        evaluation_function=lambda p: (p["x1"] + 2*p["x2"] - 7)**2 + (2*p["x1"] + p["x2"] - 5)**2,
        minimize=True,
        random_seed=2,
    )

print(f"best params: {best_parameters}")

for t in experiment.trials:
  print('---')
  print(t)
  print(experiment.trials[t].arm.parameters)
  print(experiment.trials[t].objective_mean)

## -- End pasted text --
0.1.20
[INFO 06-20 19:13:57] ax.service.utils.instantiation: Inferred value type of ParameterType.FLOAT for parameter x1. If that is not the expected value type, you can explicity specify 'value_type' ('int', 'float', 'bool' or 'str') in parameter dict.
[INFO 06-20 19:13:57] ax.service.utils.instantiation: Inferred value type of ParameterType.FLOAT for parameter x2. If that is not the expected value type, you can explicity specify 'value_type' ('int', 'float', 'bool' or 'str') in parameter dict.
[INFO 06-20 19:13:57] 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 06-20 19:13:57] ax.service.managed_loop: Started full optimization with 20 steps.
[INFO 06-20 19:13:57] ax.service.managed_loop: Running optimization trial 1...
[INFO 06-20 19:13:57] ax.service.managed_loop: Running optimization trial 2...
[INFO 06-20 19:13:57] ax.service.managed_loop: Running optimization trial 3...
[INFO 06-20 19:13:57] ax.service.managed_loop: Running optimization trial 4...
[INFO 06-20 19:13:57] ax.service.managed_loop: Running optimization trial 5...
[INFO 06-20 19:13:58] ax.service.managed_loop: Running optimization trial 6...
[INFO 06-20 19:13:58] ax.service.managed_loop: Running optimization trial 7...
[INFO 06-20 19:13:58] ax.service.managed_loop: Running optimization trial 8...
[INFO 06-20 19:13:59] ax.service.managed_loop: Running optimization trial 9...
[INFO 06-20 19:13:59] ax.service.managed_loop: Running optimization trial 10...
[INFO 06-20 19:14:00] ax.service.managed_loop: Running optimization trial 11...
[INFO 06-20 19:14:00] ax.service.managed_loop: Running optimization trial 12...
[INFO 06-20 19:14:01] ax.service.managed_loop: Running optimization trial 13...
[INFO 06-20 19:14:01] ax.service.managed_loop: Running optimization trial 14...
[INFO 06-20 19:14:02] ax.service.managed_loop: Running optimization trial 15...
[INFO 06-20 19:14:03] ax.service.managed_loop: Running optimization trial 16...
[INFO 06-20 19:14:03] ax.service.managed_loop: Running optimization trial 17...
[INFO 06-20 19:14:04] ax.service.managed_loop: Running optimization trial 18...
[INFO 06-20 19:14:05] ax.service.managed_loop: Running optimization trial 19...
[INFO 06-20 19:14:05] ax.service.managed_loop: Running optimization trial 20...

best params: {'x1': 1.4327030843905124, 'x2': 2.8136152874988056}
---
0
{'x1': 5.551732778549194, 'x2': 9.377297163009644}
539.1629722466083
---
1
{'x1': -3.933116849511862, 'x2': -7.940220162272453}
1151.8753706086577
---
2
{'x1': -5.739457290619612, 'x2': 0.686403214931488}
378.6041670452986
---
3
{'x1': 4.6975501254200935, 'x2': -1.6736418567597866}
39.3258254072101
---
4
{'x1': 2.1678781881928444, 'x2': 4.608701150864363}
34.78943015956635
---
5
{'x1': 7.36139263591334, 'x2': 1.504591000062078}
137.41475132767002
---
6
{'x1': 2.310447780652936, 'x2': 1.3904117780626084}
4.665987644465215
---
7
{'x1': -4.810674963815469, 'x2': 8.364400314275986}
63.337378518202556
---
8
{'x1': -1.5243934300008029, 'x2': 5.022754561104923}
11.47066442305054
---
9
{'x1': 9.999999999999986, 'x2': -6.252351964268424}
166.86074292618207
---
10
{'x1': -10.0, 'x2': 10.0}
234.0
---
11
{'x1': 0.6636458893856982, 'x2': 3.2849454446895585}
0.20489939790511563
---
12
{'x1': 2.629822040934295, 'x2': 2.2841183380110905}
6.509954504154573
---
13
{'x1': 0.12090649037823198, 'x2': 4.516155679782699}
4.694926278597075
---
14
{'x1': 1.4327030843905124, 'x2': 2.8136152874988056}
0.46466218161603334
---
15
{'x1': -1.5311382671850993, 'x2': 7.107052332123033}
33.20856034668283
---
16
{'x1': 3.765216423598442, 'x2': 0.6366904023760931}
13.877770510619534
---
17
{'x1': 0.9165355489969684, 'x2': 3.4964105314979896}
0.9354875919479475
---
18
{'x1': 9.999999999999993, 'x2': -2.6701034251441644}
157.50291764877784
---
19
{'x1': 1.4885066277036607, 'x2': 2.4159862924326205}
0.6161971451533785

[Balandat]

1. It's possible, there are a lot of sources of randomness (initial points, initial conditions for the optimizer for fitting theu model, heuristic for initial conditions for the acquisition function optimization, ...). We should take a look and make sure we're not missing anything with the seed.
2. Since your evaluation function only returns a scalar value, so we assume observations are noisy with unknown noise level, infer the noise level, and then return the "best predicted in-sample" trial, i.e. the trial that under the model with the inferred observation noise is the best. If your evaluation is noiseless (as in this synthetic example), try returning a tuple (value, 0.0) instead. Here 0.0 is a specified noise level of zero, and the model should return the actual "best observed" trial. For more information on this, see https://ax.dev/docs/trial-evaluation.html#evaluation-function

[kurtosis]

Thanks!
(2) sounds like a sensible default, although I think the page you linked has an incorrect description?

in the third case that the mean is for the objective and that SEM is 0

The glossary entry for SEM aligns with what you said though.

Fwiw it might be nice to have more documentation of optimize (see my other issue :) - or maybe just links to other docs. It's a great entry point to get up and running with Ax, but I wasn't sure what to read to understand what it's doing under the hood, or how to set parameters to integer, log scale, etc.

[kurtosis]

Btw, I think I was misled by the best_objectives code in section 5 of this tutorial:
https://ax.dev/tutorials/tune_cnn.html
which seems to be plotting raw observation values, as if they were assumed to be noiseless (which I think is not the case for that example?)

[Balandat]

although I think the page you linked has an incorrect description?

Ah good catch, we must have forgotten to update this part of the docs. Fixed in #606.

Fwiw it might be nice to have more documentation of optimize

Point taken - we should write up a more detailed description of what's going on under the hood.

Btw, I think I was misled by the best_objectives code in section 5 of this tutorial: https://ax.dev/tutorials/tune_cnn.html

Good point also, we should do one of the following:

1. plot the model-predicted means at the observed arms (rather than the evaluatiaon)
2. change the evaluation function to return 0.0 SEM

I don't think an option 3 (leave things as is but explain why we do it) makes a lot of sense.

[lena-kashtelyan]

Setting random_seed does not lead to reproducible output. Values only change by a small amount on subsequent runs so it's not 100% "random"- perhaps there's another RNG that feeds into one part of this?

There is not a full reproducibility guarantee that comes from setting random_seed in optimize or AxClient, unfortunately –– more information here in AxClient docstring: https://ax.dev/api/service.html#ax.service.ax_client.AxClient (there also was discussion here: #151 and a workaround suggested here: #95 (comment), although it's not a perfect one and might be hard to use in optimize where there is no individual control over each iteration). We should add this information to the docstring of optimize too.

Overall, we always recommend using the Service API for use cases that are not entirely trivial (since Loop API aka optimize offers rather limited functionality and control over the optimization). You might find it easier to use and better documented, too, @kurtosis.

[kurtosis]

Thanks @lena-kashtelyan! My two cents- optimize is a great entry point to get more folks to try out Ax. Basically anyone who's training a model with hyperparameters can drop this into their code in under an hour. The main thing I'd suggest for the tutorial is more explicit description of the options for setting the parameter space (integer, categorical, log scale) since those come up a lot. Also maybe links to the docs for the objects it returns (e.g. where can one learn what's in experiment or what exactly model is) .

[lena-kashtelyan]

This is fantastic feedback, thank you!!