Use StratifiedStandardize for per-task Y standardization in TL

Summary: Adds per-task outcome standardization to the transfer learning adapter, ensuring each task's observations are standardized independently rather than jointly. Updates the default transform pipeline to use TL-specific outcome transforms.

Differential Revision: D102197139

Author: Carl Hvarfner

ax/adapter/registry.py

@@ -139,6 +139,7 @@
 ]
 
 Y_trans: list[type[Transform]] = [Derelativize, Winsorize, BilogY, StandardizeY]
+TL_Y_trans: list[type[Transform]] = [Derelativize, Winsorize, BilogY]
 
 # Expected `List[Type[Transform]]` for 2nd anonymous parameter to
 # call `list.__add__` but got `List[Type[SearchSpaceToChoice]]`.

ax/adapter/transfer_learning/adapter.py

@@ -24,6 +24,7 @@
     GeneratorSetup,
     MBM_X_trans,
     Y_trans,
+    Y_trans,
 )
 from ax.adapter.torch import FIT_MODEL_ERROR, TorchAdapter
 from ax.adapter.transfer_learning.utils import get_joint_search_space
@@ -54,6 +55,7 @@
 from ax.utils.common.logger import get_logger
 from botorch.models.multitask import MultiTaskGP
 from botorch.models.transforms.input import InputTransform, Normalize
+from botorch.models.transforms.outcome import StratifiedStandardize
 from botorch.utils.datasets import MultiTaskDataset, SupervisedDataset
 from gpytorch.kernels.kernel import Kernel
 from pyre_extensions import assert_is_instance
@@ -846,6 +848,7 @@ def transfer_learning_generator_specs_constructor(
                     botorch_model_class=model_class,
                     model_options=botorch_model_kwargs or {},
                     input_transform_classes=input_transform_classes,
+                    outcome_transform_classes=[StratifiedStandardize],
                     input_transform_options=input_transform_options,
                     mll_options=mll_kwargs,
                     covar_module_class=covar_module_class,