diff --git a/rdagent/scenarios/data_science/proposal/exp_gen/merge.py b/rdagent/scenarios/data_science/proposal/exp_gen/merge.py
index f8e8ffe7..d28b9e71 100644
--- a/rdagent/scenarios/data_science/proposal/exp_gen/merge.py
+++ b/rdagent/scenarios/data_science/proposal/exp_gen/merge.py
@@ -119,15 +119,26 @@ class ExpGen2Hypothesis(DSProposalV2ExpGen):
         resp_dict = json.loads(response)
         return resp_dict
 
+    def get_exp_index(self, trace: DSTrace) -> int:
+        leaves: list[int] = trace.get_leaves()
+        if trace.sota_exp_to_submit is not None:
+            sota_submit_value = trace.sota_exp_to_submit.result.loc["ensemble"].iloc[0]
+            trace_scores = []
+            for i, leaf in enumerate(leaves):
+                if leaf == trace.current_selection[0]:
+                    continue
+                fb = trace.sota_experiment_fb(selection=(leaf,))
+                if fb is None:
+                    continue
+                final_score = fb[0].result.loc["ensemble"].iloc[0]
+                trace_scores.append((i, abs(final_score - sota_submit_value)))
+            if trace_scores:
+                return min(trace_scores, key=lambda item: item[1])[0]
+        return next((i for i, leaf in enumerate(leaves) if leaf != trace.current_selection[0]))
+
     def gen(self, trace: DSTrace) -> DSExperiment:
         # Ignore the selection argument and use all leaves instead.
-        leaves: list[int] = trace.get_leaves()
         sota_exp_fb = trace.sota_experiment_fb(selection=trace.current_selection)
-        exp_index = leaves[1] if trace.current_selection[0] == leaves[0] else leaves[0]
-
-        exp_to_merge_fb = trace.sota_experiment_fb(selection=(exp_index,))
-        if exp_to_merge_fb is None:
-            exp_to_merge_fb = trace.hist[exp_index]
 
         if sota_exp_fb:
             sota_exp_desc = T("scenarios.data_science.share:describe.exp").r(
@@ -139,9 +150,30 @@ class ExpGen2Hypothesis(DSProposalV2ExpGen):
             sota_exp_desc = ""
             eda_output = None
 
-        success_fb_list = trace.experiment_and_feedback_list_after_init(
-            return_type="sota", search_type="ancestors", selection=(exp_index,)
-        )
+        trace_fbs = []
+        # find the best exp to merge
+        leaves: list[int] = trace.get_leaves()
+        for leaf in leaves:
+            if leaf == trace.current_selection[0]:
+                continue
+
+            trace_fbs.append(
+                trace.experiment_and_feedback_list_after_init(
+                    return_type="sota",
+                    search_type="ancestors",
+                    selection=(leaf,),
+                )
+            )
+
+        num_to_slice = 20
+        if sum(len(fb_list) for fb_list in trace_fbs) > num_to_slice:
+            success_fb_trace_count = sum(1 for fb_list in trace_fbs if fb_list)
+            success_fb_list = [
+                fb for fb_list in trace_fbs for fb in fb_list[-(num_to_slice // success_fb_trace_count) :]
+            ]
+        else:
+            success_fb_list = [fb for fb_list in trace_fbs for fb in fb_list]
+
         if len(success_fb_list) > 0:
             exp_to_merge_fb_desc = T("scenarios.data_science.proposal.exp_gen.merge:trace").r(
                 exp_and_feedback_list=success_fb_list,
@@ -151,6 +183,11 @@ class ExpGen2Hypothesis(DSProposalV2ExpGen):
                 pipeline=DS_RD_SETTING.coder_on_whole_pipeline,
             )
         else:
+            exp_index = self.get_exp_index(trace)
+            exp_to_merge_fb = trace.sota_experiment_fb(selection=(exp_index,))
+            if exp_to_merge_fb is None:
+                exp_to_merge_fb = trace.hist[exp_index]
+
             exp_to_merge_fb_desc = T("scenarios.data_science.share:describe.feedback").r(
                 exp_and_feedback=exp_to_merge_fb,
                 heading="The feedback for the solution to be merged",
@@ -188,6 +225,7 @@ class ExpGen2Hypothesis(DSProposalV2ExpGen):
 
 
 class ExpGen2TraceAndMerge(ExpGen):
+
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.merge_exp_gen = MergeExpGen(self.scen)
@@ -396,7 +434,9 @@ class ExpGen2TraceAndMergeV3(ExpGen):
             else:
                 selection = (leaves[0],)
                 if trace.sota_exp_to_submit is not None:
-                    if trace.is_parent(trace.exp2idx(trace.sota_exp_to_submit), leaves[1]):
-                        selection = (leaves[1],)
+                    for i in range(1, len(leaves)):
+                        if trace.is_parent(trace.exp2idx(trace.sota_exp_to_submit), leaves[i]):
+                            selection = (leaves[i],)
+                            break
                 trace.set_current_selection(selection)
                 return self.merge_exp_gen.gen(trace)