[WIP] Support ScalarizedObjective children in MultiObjective

[saitcakmak]

Summary

Adds support for ScalarizedObjective as children of MultiObjective, enabling scalarized multi-objective optimization. For example, a user can optimize "0.5 * latency + 0.5 * cost, accuracy" where the first MOO objective is a scalarized combination of two metrics and the second is a plain metric.

This PR depends on the unify-objective-weight-matrix branch (PR #4943) being merged first. It builds on the unified 2D objective_weights tensor to encode scalarization directly — no separate objective_weight_matrix field is needed.

Design principle

• Standard MOO: each row of objective_weights has exactly one nonzero entry (e.g., [1, 0, 0]).
• Scalarized MOO: at least one row has multiple nonzero entries (e.g., [0.5, 0.5, 0]).
• Detection: has_scalarized_objectives(W) = (W != 0).sum(dim=1).max() > 1.

Changes by file

Core (3 files)

File	Change
ax/core/objective.py	Remove NotImplementedError blocking ScalarizedObjective in MultiObjective. Update MultiObjective.metrics to flatten component metrics from scalarized children via o.metrics instead of o.metric.
ax/core/optimization_config.py	Update objectives_by_signature dict to use obj.expression as key for ScalarizedObjective children (since they have no single .metric).
ax/core/tests/test_objective.py	Replace two negative tests with positive tests verifying MultiObjective accepts ScalarizedObjective children and .metrics returns all component metrics.

Adapter (8 files)

File	Change
ax/adapter/adapter_utils.py	extract_objective_weights: mark all component metrics of scalarized children as nonzero so subset_model keeps them. feasible_hypervolume: compute scalarized objective values from component metrics when scalarized sub-objectives exist.
ax/adapter/torch.py	Guard _untransform_objective_thresholds to skip when has_scalarized_objectives() is true (scalarized thresholds don't map back to per-metric thresholds).
ax/adapter/transforms/standardize_y.py	Add elif MultiObjective block to rescale ScalarizedObjective children's weights via _transform_scalarized_weights.
ax/adapter/transforms/stratified_standardize_y.py	Same pattern: update early-return guard and add elif MultiObjective block to rescale scalarized children's weights by strata-specific std.
ax/adapter/transforms/power_transform_y.py	Add elif MultiObjective block to raise NotImplementedError if any ScalarizedObjective child's metrics overlap with power-transformed metrics (nonlinear transform corrupts linear scalarization).
ax/adapter/transforms/objective_as_constraint.py	Handle ScalarizedObjective children by creating ScalarizedOutcomeConstraint instead of accessing .metric (which raises NotImplementedError).
ax/adapter/transforms/winsorize.py	Replace single-line minimize lookup with loop handling scalarized children by iterating metric_weights and inferring direction as minimize XOR (weight < 0).
ax/adapter/transforms/tests/test_standardize_y_transform.py	Replace negative test with positive test creating MultiObjective([ScalarizedObjective(...), Objective(...)]) and verifying weights are rescaled by sigma.

Generators (5 files)

File	Change
ax/generators/torch/utils.py	Add has_scalarized_objectives(W) helper. Add _get_scalarized_mo_objective(W) returning GenericMCMultiOutputObjective that computes samples @ W.T. Update get_botorch_objective_and_transform MOO branch to use scalarized objective when detected. Fix extract_objectives to handle 1D input.
ax/generators/torch/botorch_moo_utils.py	get_weighted_mc_objective_and_objective_thresholds: add scalarized branch using _get_scalarized_mo_objective and computing scalarized thresholds via clean_thresholds @ W.T. infer_objective_thresholds: add scalarized branch computing pred @ W.T, finding Pareto frontier, storing thresholds at first nonzero column positions.
ax/generators/torch/botorch_modular/acquisition.py	When scalarized objectives detected, replace NaN with 0 in cloned thresholds tensor before passing to input constructor.
ax/generators/torch/botorch_modular/utils.py	_objective_threshold_to_outcome_constraints: add scalarized branch where each row of objective_weights becomes one constraint row A[i] = -W_i, b[i] = -scalarized_threshold_i.

E2E test (1 file)

File	Change
ax/api/tests/test_client.py	Add test_multi_objective_with_scalarized_sub_objectives covering: both-scalarized, mixed (scalarized + plain), and weighted cases with Sobol + BoTorch trials.

Remaining TODOs

• Objective thresholds for scalarized sub-objectives: Currently, objective thresholds are always inferred via infer_objective_thresholds. There's no user-facing API to specify thresholds for scalarized sub-objectives (what metric name/expression would the ObjectiveThreshold use?). The check_objective_thresholds_match_objectives function uses obj.expression as the key for scalarized sub-objectives, but user-specified thresholds on scalarized objectives haven't been tested end-to-end.
• get_pareto_frontier_and_configs (adapter_utils.py ~line 648): This function calls extract_objective_weights (1D) instead of extract_objective_weight_matrix (2D). It may need updating for scalarized MOO to correctly compute Pareto frontiers. Currently not tested with scalarized objectives.
• hypervolume (adapter_utils.py ~line 907): Uses obj_w.nonzero().view(-1) which assumes 1D weights. May need updating for 2D weight matrix with scalarized rows.
• JSON/DB serialization: MultiObjective with ScalarizedObjective children may not serialize/deserialize correctly. The storage layer hasn't been tested with this new structure.
• Additional transform tests: The three newly updated transforms (stratified_standardize_y, power_transform_y, objective_as_constraint) pass existing tests but don't have dedicated tests for the scalarized MOO paths.
• pareto_frontier_evaluator (botorch_moo_utils.py): Uses objective_weights.shape assuming 1D in the zeros fallback (line ~159). May need updating.
• Audit ax/service/ and ax/analysis/: These modules may access MultiObjective.objectives[i].metric which would fail for ScalarizedObjective children.

Test plan

• python -m pytest ax/core/tests/test_objective.py (4 passed)
• python -m pytest ax/core/tests/test_optimization_config.py (13 passed)
• python -m pytest ax/adapter/transforms/tests/test_standardize_y_transform.py (5 passed)
• python -m pytest ax/adapter/transforms/tests/test_stratified_standardize_y_transform.py (8 passed)
• python -m pytest ax/adapter/transforms/tests/test_power_y_transform.py (3 passed)
• python -m pytest ax/adapter/transforms/tests/test_objective_as_constraint.py (27 passed)
• python -m pytest ax/generators/torch/tests/ (186 passed)
• python -m pytest ax/adapter/tests/ (144 passed)
• python -m pytest ax/api/tests/test_client.py::TestClient::test_multi_objective_with_scalarized_sub_objectives (1 passed)

[meta-codesync]

@saitcakmak has imported this pull request. If you are a Meta employee, you can view this in D94424106.