[GENERAL SUPPORT]: Updating trial parameters after get_next_trial

[davifebba]

Question

I have the following problem (using the ServiceAPI): I have time and temperature as search parameters. When I use get_next_trial(), AxClient suggests both, which is fine for temperature. However, I need to mark the trial as completed after the evaluation method, so that the suggested time is no longer valid and needs to be updated to the correct time after the evaluation method is called. For example:

next_parameters, trial_index = ax_client.get_next_trial()
next_parameters["time"] = 2.0 # for example
evaluation_method() # At this point, the time has passed. New time is new_time = 2.0 + delta_t
ax_client.complete_trial(...) # The updated time (new_time) must be informed to the client.

How to solve this problem? Using fixed_features in get_next_trial does not work because I need the updated time. So how to update the parameterization with an adjusted parameter before marking the trial as complete, using a generation strategy? In other words, how to update the arm parameters in the generation strategy before marking the trial as complete?

I'm using this generation strategy, so that Botorch takes over after a few sobol samples. If I use ax_client.attach_trial(), then it's a manual step and not part of the generation strategy.

GenerationStep(  # Initialization step
    # Which model to use for this step
    model=Models.SOBOL,
    # How many generator runs (each of which is then made a trial) to produce with this step
    num_trials=sobol_trials,
    # How many trials generated from this step must be COMPLETED before the next one
    min_trials_observed=sobol_trials,
),
GenerationStep(  # BayesOpt step
    model=Models.BOTORCH_MODULAR,
    # No limit on how many generator runs will be produced
    num_trials=botorch_trials,
    model_kwargs=self.botorch_model_kwargs,
),
])

Please provide any relevant code snippet if applicable.

No response

Code of Conduct

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

[mgrange1998]

Hi @davifebba, thanks for opening the issue.
One thing I don't understand is why "time" is part of your search space- as time goes on, more and more of the search space will become inaccessible, which doesn't really make sense for a sweep.

Some suggestions

1. Could you first write and run the code for your problem using only "temperature" as a search parameter?
2. Could you explain a bit more about the underlying problem you are solving by searching w/ "time" and "temperature" to determine whether using "time" as a parameter is appropriate.

[davifebba]

Hi @mgrange1998 !

1. Yes, it works with only temperature, and even with time using ax.client.get_next_trial(fixed_features=FixedFeatures({'time':time.time()-start_time}))
   However, the time that goes into FixedFeatures has a -delta_T to the real time I wanted to pass to the model.

2. I'm measuring the current-voltage response of a device, which degrades overtime. The idea is to use time and temperature as a parameter to build an interpolation model to predict device response as a function of these two parameters. The device not only depends on temperature, but also on its age. I'm trying to build a model that can make predictions as a function of (temperature, time) in an active learning scenario. If I use time as fixed_features, the model than can take temperature and time as independent variables and calculate the mean and uncertainty of the device response. When Ax suggests a new temperature with time as fixed feature, this would work just fine if I didn't have to wait for a temperature ramp to then measure the device response, at a different time than the one suggested.

[Balandat]

Interesting. You're essentially trying to use time as a "contextual" variable. Unfortunately, I don't think this is easily supported in the current service API, but it's something that we want to support going forward. This would require quite a bit of additional work on our internals though, so having full support for this may take some time.

That said, @sdaulton has done related things using a TimeAsFeature transform. This is not the exactly what you'd need here since this uses the trial start and end times rather than the age of the device. What you'd need is some other transform that converts the "time" metric into an addiitonal parameter for the model.

[davifebba]

Exactly, time in this case would be a "contextual" variable. Since we cannot control time, the model would suggest only temperature points, but also use time to make predictions. In my case, "age" means the elapsed time since the start of the device characterization. Would it be possible to implement something using the Developer API, and informing the model of the correct time after a new trial has been suggested?

[Balandat]

Would it be possible to implement something using the Developer API, and informing the model of the correct time after a new trial has been suggested?

Yes, that should be possible. We can implement a new Transform that (i) add the new time parameter as a parameter to the (transformed) search space and (ii) grabs the specified metric from the observation features and injects it as the parameter value instead. This will be very similar to the TimeAsFeature transform. We would then have to specify this as a transform to use in a custom GenerationStrategy setup.

[davifebba]

@Balandat, any guidance or example on how to implement that would be very appreciated!

[sdaulton]

Hi @davifebba! It seems like using TimeAsFeature would work here. Although this uses the start time/duration as a contextual variables, the start time/duration should just be the device age with some additive offset, right?

For a given trial, you can report the start/end time for the evaluation to Ax by passing "start_time" and "end_time" via the metadata dictionary in AxClient.complete_trial (https://github.com/facebook/Ax/blob/main/ax/service/ax_client.py#L770). If you then include the TimeAsFeature transform in the BOTORCH_MODULAR GenerationStep (e.g. by passing model_kwargs={"transforms": Cont_X_trans+TimeAsFeature + Y_trans}), then start time and duration will be added as contextual features to the model.

Does that address the problem or am I misunderstanding?

[davifebba]

Hi @sdaulton! Thanks for the help! So, what I want to accomplish in the end is to use temperature and time as independent variables to predict device performance. In this case, time is any real time instant in the future. Do you mean to implement something like this?

optimization_parameters = [
    {"name": "temperature", 
     "type": "range", 
     "bounds": [300, 550], 
     "value_type": "int"
     }
]
model_kwargs={"transforms": Cont_X_trans+[TimeAsFeature] + Y_trans}

start_time = time.time()
BO loop:
   next_parameters, trial_index = self.optimizer.ax_client.get_next_trial()
   (time consuming steps)
   elapsed_time = (time.time()-start_time)/3600 # In hours
   complete_trial(...metadata={"start_time": start_time, "end_time": elapsed_time)

I tried that and got this error, which I don't know what it means:
ax.exceptions.core.UserInputError: Upper bound of start_time must be strictly larger than lower.Got: (1.732642, 1.732642).

Also, would this enable the model to use time as a predictive variable?

[sdaulton]

Hi @davifebba, sorry for the delay. That looks good, but "end_time" should be the current time, not the elapsed time since the start.

[davifebba]

Hi @sdaulton ! The code below still gives me an error:

\ax\core\parameter.py", line 304, in _validate_range_param
    raise UserInputError(
ax.exceptions.core.UserInputError: Upper bound of start_time must be strictly larger than lower.Got: (1.740172, 1.740172).

optimization_parameters = [
    {"name": "temperature", 
     "type": "range", 
     "bounds": [300, 550], 
     "value_type": "int"
     }
]
model_kwargs={"transforms": Cont_X_trans+[TimeAsFeature] + Y_trans}

start_time = time.time()
BO loop:
   next_parameters, trial_index = self.optimizer.ax_client.get_next_trial()
   (time consuming steps)
   end_time = time.time()
   complete_trial(...metadata={"start_time": start_time, "end_time": end_time})

Any suggestions on how to fix this error?