[Bug]: ChoiceParameters out-of-design not included in training data

[benrabe93]

What happened?

Hey.

I've noticed that completed trials with ChoiceParameter-values out-of-design space are NOT used for fitting the GP / candidate generation

I've provided a minimum working example below. Note that the GP uncertainty is very high at the outer points x = 0.0 and x = 10.0 and the constant mean, even though completed trials exist nearby. In fact, the GP mean & uncertainty is the same if those completed trials nearby were not provided at all (check by commenting out those points).

And I assume that this issue also applies to ChoiceParameters of type string.

I'd very much appreciate if you can show me a work-around or fix the potential bug.
Thanks!

Please provide a minimal, reproducible example of the unexpected behavior.

from ax.api.client import Client
from ax.api.configs import ChoiceParameterConfig

def parabola(x):
    return x**2

choices = [0.0, 5.0, 10.0]

client = Client()
client.configure_experiment(parameters=[ChoiceParameterConfig(name="x", values=sorted(choices), parameter_type="float", is_ordered=True)])
client.configure_optimization(objective="-parabola")
client.configure_generation_strategy(initialize_with_center=False, initialization_budget=1)

preexisting_data = [
    {"x": -0.01},
    {"x": 0.01},
    {"x": 5.0},
    {"x": 10.01},
    {"x": 9.99},
]
for params in preexisting_data:
    trial_index = client.attach_trial(parameters=params)
    client.complete_trial(trial_index=trial_index, raw_data={"parabola": parabola(params["x"])})

# Generate trial using MBM
client.get_next_trials(max_trials=1)

# Compute predictions
predict_data = [{"x": c} for c in choices]
preds = client.predict(predict_data)
for i in range(len(choices)):
    print(f"{predict_data[i]}: {preds[i]}")

Output

{'x': 0.0}: {'parabola': (np.float64(25.0), np.float64(0.9987194172385941))}
{'x': 5.0}: {'parabola': (np.float64(25.0), np.float64(0.08167443742157637))}
{'x': 10.0}: {'parabola': (np.float64(25.0), np.float64(0.9987194172385941))}

Ax Version

1.2.4

Python Version

3.12.13

Operating System

Ubuntu 26

(Optional) Describe any potential fixes you've considered to the issue outlined above.

No response

Pull Request

None

Code of Conduct

• I agree to follow Ax's Code of Conduct

[benrabe93]

Another related issue: If I run above code example without the datapoint at x = 5:

preexisting_data = [
    {"x": -0.01},
    {"x": 0.01},
    # {"x": 5.0},
    {"x": 10.01},
    {"x": 9.99},
]

then a DataRequiredError throws: DataRequiredError: 'MapKeyToFloat' transform requires non-empty data.

[saitcakmak]

Thanks for flagging @benrabe93. This is something that wasn't supported due to legacy transform reasons, but those blockers are no longer in place. #4217 was created to address this but looks like it fell through the cracks. I'll look into reviving & landing that.