[GENERAL SUPPORT]: MultiTask with Choice Parameter using Ax Service API

[ritalyu17]

Question

I have a function that takes both range and choice parameters as inputs, and I want to perform multi-task on it. I've included "TaskChoiceToIntTaskChoice, UnitX, ChoiceToNumericChoice" as the transforms, but I still get the error:

ValueError: This generator operates on [0,1]^d. Please make use of the UnitX transform in the ModelBridge, and ensure task features are fixed.

My question is can i specify which transform to apply to which parameter when using the Ax Service API? Or is Ax (not the service API) recommended if performing multi-task with range and choice parameters?

Any help is appreciated!

Please provide any relevant code snippet if applicable.

import numpy as np
from ax.core.observation import ObservationFeatures
from ax.modelbridge.generation_strategy import GenerationStep, GenerationStrategy
from ax.modelbridge.registry import Models
from ax.modelbridge.transforms.task_encode import TaskEncode, TaskChoiceToIntTaskChoice
from ax.modelbridge.transforms.unit_x import UnitX
from ax.modelbridge.transforms.choice_encode import ChoiceEncode, ChoiceToNumericChoice
from ax.service.ax_client import AxClient, ObjectiveProperties

# Branin function with an additional c1 parameter
def branin(x1, x2, c1):
    a = 1.0
    b = 5.1 / (4.0 * np.pi**2)
    c = 5.0 / np.pi
    r = 6.0
    s = 10.0
    t = 1.0 / (8.0 * np.pi)
    
    # Adjust behavior based on c1 string value
    if c1 == "X":
        adjustment = 0
    elif c1 == "Y":
        adjustment = 2
    elif c1 == "Z":
        adjustment = 5
    
    return a * (x2 - b * x1**2 + c * x1 - r) ** 2 + s * (1 - t) * np.cos(x1) + s + adjustment

# Shifted and inverted Branin function with an additional c1 parameter (string choices)
def shifted_inverted_branin(x1, x2, c1):
    return -branin(x1 + 2.5, x2 + 2.5, c1) + 300

transforms = [TaskChoiceToIntTaskChoice, UnitX, ChoiceToNumericChoice]

gs = GenerationStrategy(
    name="MultiTaskOp",
    steps=[
        GenerationStep(
            model=Models.SOBOL,
            num_trials=10,
            model_kwargs={"deduplicate": True, "transforms": transforms},
        ),
        GenerationStep(
            model=Models.BOTORCH_MODULAR,
            num_trials=-1,
            model_kwargs={"transforms": transforms},
        ),
    ],
)

ax_client = AxClient(generation_strategy=gs, random_seed=42, verbose_logging=False)

ax_client.create_experiment(
    name="MultiTaskOp",
    parameters=[
        {
            "name": "x1",
            "type": "range",
            "value_type": "float",
            "bounds": [-5.0, 10.0],
        },
        {
            "name": "x2",
            "type": "range",
            "value_type": "float",
            "bounds": [0.0, 15.0],
        },
        {
            "name": "c1",
            "type": "choice",
            "values": ["X", "Y", "Z"],  # string values for c1
        },
        {
            "name": "Task",
            "type": "choice",
            "values": ["A", "B"],
            "is_task": True,
            "target_value": "B",
        },
    ],
    objectives={"Objective": ObjectiveProperties(minimize=False)},
)

for i in range(40):
    p, trial_index = ax_client.get_next_trial(
        fixed_features=ObservationFeatures({"Task": "A" if i % 2 else "B"})
    )

    # Call branin or shifted_inverted_branin with x1, x2, and c1
    if p["Task"] == "A":
        u = branin(p["x1"], p["x2"], p["c1"])
    else:
        u = shifted_inverted_branin(p["x1"], p["x2"], p["c1"])

    ax_client.complete_trial(trial_index=trial_index, raw_data={"Objective": u})

Code of Conduct

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

[danielcohenlive]

Hi @ritalyu17, thanks for the question. Our tutorial is a little out of date here (curious on what you were working off of). I think you should be using Specified_Task_ST_MTGP_trans for your transformations. So it would look like:

from ax.modelbridge.registry import Specified_Task_ST_MTGP_trans
...
gs = GenerationStrategy(
    name="MultiTaskOp",
    steps=[
        GenerationStep(
            model=Models.SOBOL,
            num_trials=10,
            model_kwargs={"deduplicate": True, "transforms": Specified_Task_ST_MTGP_trans},
        ),
        GenerationStep(
            model=Models.BOTORCH_MODULAR,
            num_trials=-1,
            model_kwargs={"transforms": Specified_Task_ST_MTGP_trans},
        ),
    ],
)
...

I tried and it appears to be working and uses TaskChoiceToIntTaskChoice, but cc @Balandat to double check me on this.

[ritalyu17]

Hi @danielcohenlive, thank for the information. Using Specified_Task_ST_MTGP_trans works.

I am working on an optimization problem that takes number and strings as input. I use Branin function here as a prototype to make sure multi-task with number & string inputs works. I was mostly following the set up here, issue #2546 , and browsing other similar questions #979 , #2475 , and #2514 .

Marking this as close, will re-open if I encounter further issues.