Optimization trace to return noiseless when add_custom_noise is used

ax/benchmark/benchmark.py

@@ -190,7 +190,9 @@ def get_oracle_experiment_from_params(
         optimization_config=problem.optimization_config,
     )
 
-    runner = BenchmarkRunner(test_function=problem.test_function, noise_std=0.0)
+    # Ensure noiseless evaluation by removing any custom noise function
+    test_function = replace(problem.test_function, add_custom_noise=None)
+    runner = BenchmarkRunner(test_function=test_function, noise_std=0.0)
 
     # Silence INFO logs from ax.core.experiment that state "Attached custom
     # parameterizations"

ax/benchmark/tests/test_benchmark.py

@@ -995,6 +995,21 @@ def test_get_oracle_experiment_from_params(self) -> None:
                 problem=problem, dict_of_dict_of_params={0: {}}
             )
 
+        with self.subTest("custom noise is ignored for noiseless oracle"):
+            problem_with_custom_noise = dataclasses.replace(
+                problem,
+                test_function=dataclasses.replace(
+                    problem.test_function,
+                    add_custom_noise=lambda df, *args: df.assign(mean=999.0, sem=0.0),
+                ),
+            )
+            oracle_exp = get_oracle_experiment_from_params(
+                problem=problem_with_custom_noise,
+                dict_of_dict_of_params={0: {"0": near_opt_params}},
+            )
+            df = oracle_exp.fetch_data().df
+            self.assertAlmostEqual(df["mean"].iloc[0], Branin._optimal_value, places=5)
+
     def _test_multi_fidelity_or_multi_task(
         self, fidelity_or_task: Literal["fidelity", "task"]
     ) -> None:

ax/core/batch_trial.py

@@ -508,6 +508,10 @@ def clone_to(
         """Clone the trial and attach it to a specified experiment.
         If None provided, attach it to the current experiment.
 
+        When cloning to a different experiment, the trial's index is preserved.
+        When cloning to the same experiment, a new index is auto-assigned to
+        avoid collision.
+
         Args:
             experiment: The experiment to which the cloned trial will belong.
                 If unspecified, uses the current experiment.
@@ -517,15 +521,22 @@ def clone_to(
         Returns:
             A new instance of the trial.
         """
-        use_old_experiment = experiment is None
+        cloning_to_same_experiment = (
+            experiment is None or experiment is self._experiment
+        )
         experiment = self._experiment if experiment is None else experiment
-        new_trial = experiment.new_batch_trial(
+        # Preserve trial index when cloning to a different experiment;
+        # auto-assign new index when cloning to the same experiment.
+        new_trial = BatchTrial(
+            experiment=experiment,
             trial_type=None if clear_trial_type else self._trial_type,
             ttl_seconds=self._ttl_seconds,
             generator_runs=[
-                gr if use_old_experiment else gr.clone() for gr in self.generator_runs
+                gr if cloning_to_same_experiment else gr.clone()
+                for gr in self.generator_runs
             ],
             should_add_status_quo_arm=include_sq and self.should_add_status_quo_arm,
+            index=None if cloning_to_same_experiment else self.index,
         )
         self._update_trial_attrs_on_clone(new_trial=new_trial)
         return new_trial

ax/core/tests/test_experiment.py

@@ -1423,8 +1423,10 @@ def test_clone_with(self) -> None:
         experiment._data_by_trial
         cloned_experiment = experiment.clone_with(trial_indices=[1])
         self.assertEqual(len(cloned_experiment.trials), 1)
-        cloned_df = cloned_experiment.lookup_data_for_trial(0).df
-        self.assertEqual(cloned_df["trial_index"].iloc[0], 0)
+        # Trial index is preserved when cloning to a different experiment.
+        self.assertIn(1, cloned_experiment.trials)
+        cloned_df = cloned_experiment.lookup_data_for_trial(1).df
+        self.assertEqual(cloned_df["trial_index"].iloc[0], 1)
 
         # Clone with MapData.
         experiment = get_test_map_data_experiment(
@@ -2506,3 +2508,38 @@ def test_extract_relevant_trials(self) -> None:
             )
             self.assertEqual(len(trials), 1)
             self.assertEqual(trials[0].index, 0)
+
+    def test_clone_with_preserves_trial_indices(self) -> None:
+        """Test that clone_with preserves original trial indices when cloning
+        to a different experiment."""
+        experiment = get_branin_experiment(
+            with_batch=True,
+            with_completed_trial=True,
+            num_batch_trial=10,
+            with_completed_batch=True,
+        )
+
+        # Clone a sparse subset of trials (indices 2, 5, 8)
+        subset_indices = [2, 5, 8]
+        cloned_experiment = experiment.clone_with(trial_indices=subset_indices)
+
+        # Verify trial indices are preserved exactly
+        self.assertEqual(set(cloned_experiment.trials.keys()), set(subset_indices))
+
+        # Verify each trial's index property matches
+        for trial_index in subset_indices:
+            self.assertEqual(cloned_experiment.trials[trial_index].index, trial_index)
+            # Verify status is preserved
+            original_trial = experiment.trials[trial_index]
+            cloned_trial = cloned_experiment.trials[trial_index]
+            self.assertEqual(cloned_trial.status, original_trial.status)
+
+        # Verify data trial indices are also preserved
+        for trial_index in subset_indices:
+            if trial_index in cloned_experiment._data_by_trial:
+                trial_data = cloned_experiment.lookup_data_for_trial(trial_index)
+                if not trial_data.df.empty:
+                    self.assertTrue(
+                        all(trial_data.df["trial_index"] == trial_index),
+                        f"Data trial_index mismatch for trial {trial_index}",
+                    )

ax/core/trial.py

@@ -301,6 +301,10 @@ def clone_to(
         """Clone the trial and attach it to the specified experiment.
         If no experiment is provided, the original experiment will be used.
 
+        When cloning to a different experiment, the trial's index is preserved.
+        When cloning to the same experiment, a new index is auto-assigned to
+        avoid collision.
+
         Args:
             experiment: The experiment to which the cloned trial will belong.
                 If unspecified, uses the current experiment.
@@ -310,10 +314,17 @@ def clone_to(
         Returns:
             A new instance of the trial.
         """
+        cloning_to_same_experiment = (
+            experiment is None or experiment is self._experiment
+        )
         experiment = self._experiment if experiment is None else experiment
-        new_trial = experiment.new_trial(
+        # Preserve trial index when cloning to a different experiment;
+        # auto-assign new index when cloning to the same experiment.
+        new_trial = Trial(
+            experiment=experiment,
             ttl_seconds=self.ttl_seconds,
             trial_type=None if clear_trial_type else self.trial_type,
+            index=None if cloning_to_same_experiment else self.index,
         )
         if self.generator_run is not None:
             new_trial.add_generator_run(self.generator_run.clone())