feat: merge selectively (#888)

* chore: avoid incorporate changes best as sota merge hypothesis fix: max_retrieve_num after decision chore: select last experiments and feedbacks * add the set_current_selection before the exp_gen when merging add trace.NEW_ROOT fix: no scen_prob_multiplier fix: use regex with timeout chore: hypothesis_rank with selected_idx chore: define is_parent in proposal chore: rename collect_all_ancestors to get_parent_exps --------- Co-authored-by: you-n-g <you-n-g@users.noreply.github.com> Co-authored-by: Young <afe.young@gmail.com> Co-authored-by: Xu <v-xuminrui@microsoft.com> Co-authored-by: Xu Yang <xuyang1@microsoft.com>
2026-07-28 16:07:46 +00:00 · 2025-06-04 14:46:47 +08:00
parent 4aae5b6854
commit bb71c18080
7 changed files with 409 additions and 152 deletions
@@ -3,15 +3,13 @@
 from __future__ import annotations

 from abc import ABC, abstractmethod
-from typing import Generic, TypeVar
+from typing import Generic, List, Tuple, TypeVar

 from rdagent.core.evaluation import Feedback
 from rdagent.core.experiment import ASpecificExp, Experiment
 from rdagent.core.knowledge_base import KnowledgeBase
 from rdagent.core.scenario import Scenario

-# class data_ana: XXX
-

 class Hypothesis:
    """
@@ -105,6 +103,7 @@ ASpecificKB = TypeVar("ASpecificKB", bound=KnowledgeBase)

 class Trace(Generic[ASpecificScen, ASpecificKB]):
    NodeType = tuple[Experiment, ExperimentFeedback]  # Define NodeType as a new type representing the tuple
+    NEW_ROOT: Tuple = ()

    def __init__(self, scen: ASpecificScen, knowledge_base: ASpecificKB | None = None) -> None:
        self.scen: ASpecificScen = scen
@@ -116,6 +115,7 @@ class Trace(Generic[ASpecificScen, ASpecificKB]):

        # TODO: self.hist is 2-tuple now, remove hypothesis from it, change old code for this later.
        self.knowledge_base: ASpecificKB | None = knowledge_base
+        self.current_selection: tuple[int, ...] = (-1,)

    def get_sota_hypothesis_and_experiment(self) -> tuple[Hypothesis | None, Experiment | None]:
        """Access the last experiment result, sub-task, and the corresponding hypothesis."""
@@ -126,6 +126,77 @@ class Trace(Generic[ASpecificScen, ASpecificKB]):

        return None, None

+    def is_selection_new_tree(self, selection: tuple[int, ...] | None = None) -> bool:
+        """
+        Check if the current trace is a new tree.
+        - selection maybe (-1,) when the dag_parent is empty.
+        """
+        if selection is None:
+            selection = self.get_current_selection()
+
+        return selection == self.NEW_ROOT or len(self.dag_parent) == 0
+
+    def get_current_selection(self) -> tuple[int, ...]:
+        return self.current_selection
+
+    def set_current_selection(self, selection: tuple[int, ...]) -> None:
+        self.current_selection = selection
+
+    def get_parent_exps(
+        self,
+        selection: tuple[int, ...] | None = None,
+    ) -> list[Trace.NodeType]:
+        """
+        Collect all ancestors of the given selection.
+        The return list follows the order of [root->...->parent->current_node].
+        """
+        if selection is None:
+            selection = self.get_current_selection()
+
+        if self.is_selection_new_tree(selection):
+            return []
+
+        return [self.hist[i] for i in self.get_parents(selection[0])]
+
+    def exp2idx(self, exp: Experiment | List[Experiment]) -> int | List[int] | None:
+        if isinstance(exp, list):
+            exps: List[Experiment] = exp
+
+            # keep the order
+            exp_to_index: dict[Experiment, int] = {_exp: i for i, (_exp, _) in enumerate(self.hist)}
+            return [exp_to_index[_exp] for _exp in exps]
+        else:
+            for i, (_exp, _) in enumerate(self.hist):
+                if _exp == exp:
+                    return i
+        return None
+
+    def idx2exp(self, idx: int | List[int]) -> Experiment | List[Experiment]:
+        if isinstance(idx, list):
+            idxs: List[int] = idx
+            return [self.hist[_idx][0] for _idx in idxs]
+        else:
+            return self.hist[idx][0]
+
+    def is_parent(self, parent_idx: int, child_idx: int) -> bool:
+        ancestors = self.get_parents(child_idx)
+        return parent_idx in ancestors
+
+    def get_parents(self, child_idx: int) -> List[int]:
+        if self.is_selection_new_tree((child_idx,)):
+            return []
+
+        ancestors: List[int] = []
+        curr = child_idx
+        while True:
+            ancestors.insert(0, curr)
+            parent_tuple = self.dag_parent[curr]
+            if not parent_tuple or parent_tuple[0] == curr:
+                break
+            curr = parent_tuple[0]
+
+        return ancestors
+

 class CheckpointSelector:
    """
@@ -1,5 +1,5 @@
 from abc import abstractmethod
-from typing import Literal
+from typing import List, Literal

 from rdagent.app.data_science.conf import DS_RD_SETTING
 from rdagent.core.evolving_framework import KnowledgeBase
@@ -61,8 +61,6 @@ class DSTrace(Trace[DataScienceScen, KnowledgeBase]):

        self.knowledge_base = knowledge_base

-        self.current_selection: tuple[int, ...] = (-1,)
-
        self.sota_exp_to_submit: DSExperiment | None = None  # grab the global best exp to submit

    COMPLETE_ORDER = ("DataLoadSpec", "FeatureEng", "Model", "Ensemble", "Workflow")
@@ -70,12 +68,6 @@ class DSTrace(Trace[DataScienceScen, KnowledgeBase]):
    def set_sota_exp_to_submit(self, exp: DSExperiment) -> None:
        self.sota_exp_to_submit = exp

-    def get_current_selection(self) -> tuple[int, ...]:
-        return self.current_selection
-
-    def set_current_selection(self, selection: tuple[int, ...]) -> None:
-        self.current_selection = selection
-
    @property
    def sub_trace_count(self) -> int:
        return len(self.get_leaves())
@@ -144,50 +136,11 @@ class DSTrace(Trace[DataScienceScen, KnowledgeBase]):
            return self.hist

        elif search_type == "ancestors":
-
-            if selection is None:
-                selection = self.get_current_selection()
-
-            if len(selection) == 0:
-                # selection is (), which means we switch to a new trace
-                return []
-
-            return self.collect_all_ancestors(selection)
+            return self.get_parent_exps(selection)

        else:
            raise ValueError(f"Invalid search type: {search_type}")

-    def collect_all_ancestors(
-        self,
-        selection: tuple[int, ...] | None = None,
-    ) -> list[tuple[DSExperiment, ExperimentFeedback]]:
-        """
-        Collect all ancestors of the given selection.
-        The return list follows the order of [root->...->parent->current_node].
-        """
-        if selection is None:
-            selection = self.get_current_selection()
-
-        if len(self.dag_parent) == 0:
-            return []
-
-        else:
-            all_ancestors = []
-
-            # start from the latest selection
-            current_node_idx = selection[0]
-
-            # add the current node to the list
-            all_ancestors.insert(0, self.hist[current_node_idx])
-
-            parent_idx = self.dag_parent[current_node_idx]
-
-            while len(parent_idx) > 0:
-                all_ancestors.insert(0, self.hist[parent_idx[0]])
-                parent_idx = self.dag_parent[parent_idx[0]]
-
-        return all_ancestors
-
    def next_incomplete_component(
        self,
        search_type: Literal["all", "ancestors"] = "ancestors",
@@ -226,10 +179,6 @@ class DSTrace(Trace[DataScienceScen, KnowledgeBase]):
        Retrieve a list of experiments and feedbacks based on the return_type.
        """
        search_list = self.retrieve_search_list(search_type, selection=selection)
-        if max_retrieve_num is not None and len(search_list) > 0:
-            retrieve_num = min(max_retrieve_num, len(search_list))
-            search_list = search_list[:retrieve_num]
-
        final_component = self.COMPLETE_ORDER[-1]
        has_final_component = True if DS_RD_SETTING.coder_on_whole_pipeline else False
        SOTA_exp_and_feedback_list = []
@@ -243,6 +192,13 @@ class DSTrace(Trace[DataScienceScen, KnowledgeBase]):
                    failed_exp_and_feedback_list.append((exp, fb))
            if exp.hypothesis.component == final_component and fb:
                has_final_component = True
+        if max_retrieve_num is not None and (SOTA_exp_and_feedback_list or failed_exp_and_feedback_list):
+            SOTA_exp_and_feedback_list = SOTA_exp_and_feedback_list[
+                -min(max_retrieve_num, len(SOTA_exp_and_feedback_list)) :
+            ]
+            failed_exp_and_feedback_list = failed_exp_and_feedback_list[
+                -min(max_retrieve_num, len(failed_exp_and_feedback_list)) :
+            ]
        if return_type == "all":
            return SOTA_exp_and_feedback_list + failed_exp_and_feedback_list
        elif return_type == "failed":
@@ -56,7 +56,7 @@ class LimitTimeCKPSelector(CheckpointSelector):

        Returns:
            (-1,): Continue with the current latest trial
-            (): Start a new sub-trace if max trace limit not reached
+            trace.NEW_ROOT: Start a new sub-trace if max trace limit not reached
        """

        if self.time_limit_pre_trace is None:
@@ -69,8 +69,8 @@ class LimitTimeCKPSelector(CheckpointSelector):
            logger.info(f"Starting initial sub-trace {trace.sub_trace_count} at {current_time}")
            return (-1,)  # Continue with latest trial for new sub-trace

-        # Calculate elapsed time for current sub-trace
-        elapsed_time = current_time - self.sub_trace_start_times[trace.sub_trace_count - 1]
+        # Calculate elapsed time for current sub-trace, Trace count may be larger than MAX_TRACE_NUM druing merge process
+        elapsed_time = current_time - self.sub_trace_start_times[min(trace.sub_trace_count, self.MAX_TRACE_NUM) - 1]

        if elapsed_time < self.time_limit_pre_trace:
            # Continue with current sub-trace
@@ -94,7 +94,7 @@ class LimitTimeCKPSelector(CheckpointSelector):
                f"Elapsed time {elapsed_time} exceeds time limit {self.time_limit_pre_trace}, jump to a new sub-trace"
            )
            logger.info(f"current sub-trace count: {trace.sub_trace_count}")
-            return tuple()  # Empty tuple signals starting a new sub-trace
+            return trace.NEW_ROOT  # Empty tuple signals starting a new sub-trace


 class SOTAJumpCKPSelector(CheckpointSelector):
@@ -140,7 +140,7 @@ class SOTAJumpCKPSelector(CheckpointSelector):
                    f"SOTA count {sota_count} is below threshold {self.SOTA_COUNT_THRESHOLD}, jump to a new sub-trace"
                )
                logger.info(f"current sub-trace count: {trace.sub_trace_count}")
-                return ()
+                return trace.NEW_ROOT
            else:
                logger.info(
                    f"SOTA count {sota_count} is above threshold {self.SOTA_COUNT_THRESHOLD}, continue the current latest trial"
@@ -201,7 +201,7 @@ class BackJumpCKPSelector(CheckpointSelector):
                    logger.info(
                        f"SOTA count {sota_count} is below threshold {self.SOTA_COUNT_THRESHOLD}, jump a new sub-trace"
                    )
-                    return ()  # reboot a new sub-trace
+                    return trace.NEW_ROOT  # reboot a new sub-trace
                else:
                    logger.info(
                        f"SOTA count {sota_count} is below threshold {self.SOTA_COUNT_THRESHOLD}, jump back to the last second SOTA in hist (may not in current sub-trace)"
@@ -227,7 +227,7 @@ class BackJumpCKPSelector(CheckpointSelector):
                            f"SOTA count {sota_count} is below threshold {self.SOTA_COUNT_THRESHOLD}, jump a new sub-trace"
                        )
                        logger.info(f"current sub-trace count: {trace.sub_trace_count}")
-                        return ()  # reboot a new sub-trace
+                        return trace.NEW_ROOT  # reboot a new sub-trace

            else:
                logger.info(
@@ -1,16 +1,21 @@
 """Merge the version in different traces"""

+import json
 from datetime import timedelta
+from typing import Dict, Tuple

 from rdagent.app.data_science.conf import DS_RD_SETTING
 from rdagent.components.coder.data_science.pipeline.exp import PipelineTask
 from rdagent.core.proposal import ExpGen
 from rdagent.log import rdagent_logger as logger
 from rdagent.log.timer import RD_Agent_TIMER_wrapper, RDAgentTimer
+from rdagent.oai.llm_utils import APIBackend
 from rdagent.scenarios.data_science.experiment.experiment import DSExperiment
 from rdagent.scenarios.data_science.loop import DataScienceRDLoop
 from rdagent.scenarios.data_science.proposal.exp_gen.base import DSHypothesis, DSTrace
+from rdagent.scenarios.data_science.proposal.exp_gen.proposal import DSProposalV2ExpGen
 from rdagent.utils.agent.tpl import T
+from rdagent.utils.workflow import wait_retry


 class MergeExpGen(ExpGen):
@@ -82,7 +87,106 @@ class MergeExpGen(ExpGen):
        return exp


-# dual-target version
+class ExpGen2Hypothesis(DSProposalV2ExpGen):
+    @wait_retry(retry_n=5)
+    def hypothesis_gen(
+        self,
+        component_desc: str,
+        sota_exp_desc: str,
+        enable_idea_pool: bool,
+        pipeline: bool = True,
+        exp_feedback_list_desc: str = "",
+        scenario_desc: str = "",
+        problems: dict = {},
+    ) -> Dict:
+        sys_prompt = T(".merge:hypothesis_gen.system").r(
+            component_desc=component_desc,
+            hypothesis_output_format=T(".prompts_v2:output_format.hypothesis").r(
+                pipeline=pipeline, enable_idea_pool=enable_idea_pool
+            ),
+            pipeline=pipeline,
+        )
+        user_prompt = T(".merge:hypothesis_gen.user").r(
+            exp_and_feedback_list_desc=exp_feedback_list_desc,
+            sota_exp_desc=sota_exp_desc,
+        )
+        response = APIBackend().build_messages_and_create_chat_completion(
+            user_prompt=user_prompt,
+            system_prompt=sys_prompt,
+            json_mode=True,
+            json_target_type=Dict[str, Dict[str, str | Dict[str, str | int]]],
+        )
+        resp_dict = json.loads(response)
+        return resp_dict
+
+    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(
+                exp=sota_exp_fb[0],
+                heading="Best previous exploration of the scenario",
+            )
+            eda_output = sota_exp_fb[0].experiment_workspace.file_dict.get("EDA.md", None)
+        else:
+            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,)
+        )
+        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,
+                type="success",
+                heading="Successful iterations:",
+                success_trial_desc="These trials are the steps or changes that led to the success of the solution to be merged",
+                pipeline=DS_RD_SETTING.coder_on_whole_pipeline,
+            )
+        else:
+            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",
+            )
+
+        component_desc = T("scenarios.data_science.share:component_description_in_pipeline").r()
+        hypothesis_dict = self.hypothesis_gen(
+            component_desc=component_desc,
+            exp_feedback_list_desc=exp_to_merge_fb_desc,
+            sota_exp_desc=sota_exp_desc,
+            enable_idea_pool=DS_RD_SETTING.enable_knowledge_base,
+            pipeline=DS_RD_SETTING.coder_on_whole_pipeline,
+        )
+
+        all_problems = {}
+        pickled_problem_name, new_hypothesis = self.hypothesis_rank(
+            hypothesis_dict=hypothesis_dict,
+            problem_dict=all_problems,
+            selected_idx=0,
+        )
+        if DS_RD_SETTING.enable_knowledge_base:
+            trace.knowledge_base.update_pickled_problem(all_problems, pickled_problem_name)
+
+        scenario_desc = trace.scen.get_scenario_all_desc(eda_output=eda_output)
+
+        return self.task_gen(
+            component_desc=component_desc,
+            scenario_desc=scenario_desc,
+            sota_exp_desc=sota_exp_desc,
+            sota_exp=sota_exp_fb[0] if sota_exp_fb else None,
+            hypothesis=new_hypothesis,
+            pipeline=DS_RD_SETTING.coder_on_whole_pipeline,
+            failed_exp_feedback_list_desc="",
+        )
+
+
 class ExpGen2TraceAndMerge(ExpGen):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
@@ -98,7 +202,7 @@ class ExpGen2TraceAndMerge(ExpGen):
        if timer.remain_time_duration >= timedelta(hours=DS_RD_SETTING.merge_hours):
            leaves: list[int] = trace.get_leaves()
            if len(leaves) < 2:
-                selection = tuple()  # create new trace
+                selection = trace.NEW_ROOT  # create new trace
            else:
                selection = (
                    leaves[0],
@@ -153,7 +257,6 @@ class MergeExpGen_MultiTrace(ExpGen):
                selection=(leaves[i],),
            )
            if len(success_fb_list) > 0:
-
                exp_to_merge_fb_desc = T("scenarios.data_science.share:describe.trace").r(
                    exp_and_feedback_list=success_fb_list,
                    type="success",
@@ -199,7 +302,6 @@ class ExpGen2TraceAndMergeV2(ExpGen):
        self.exp_gen = DataScienceRDLoop._get_exp_gen(
            "rdagent.scenarios.data_science.proposal.exp_gen.DSExpGen", self.scen
        )
-        self.MAX_TRACE_NUM = DS_RD_SETTING.max_trace_num  # maximum number of traces to grow before merging
        self.flag_start_merge = False

    def gen(self, trace: DSTrace) -> DSExperiment:
@@ -231,9 +333,51 @@ class ExpGen2TraceAndMergeV2(ExpGen):
            else:
                if not self.flag_start_merge:  # root node of the merge trace
                    self.flag_start_merge = True
-                    trace.set_current_selection(tuple())
+                    trace.set_current_selection(trace.NEW_ROOT)
                    return self.merge_exp_gen.gen(trace)
                else:
                    # return self.merge_exp_gen.gen(trace)
                    trace.set_current_selection(selection=(-1,))
                    return self.exp_gen.gen(trace)  # continue the last trace, to polish the merged solution
+
+
+class ExpGen2TraceAndMergeV3(ExpGen):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.merge_exp_gen = ExpGen2Hypothesis(self.scen)
+        self.exp_gen = DataScienceRDLoop._get_exp_gen(
+            "rdagent.scenarios.data_science.proposal.exp_gen.DSExpGen", self.scen
+        )
+
+    def gen(self, trace: DSTrace) -> DSExperiment:
+        timer: RDAgentTimer = RD_Agent_TIMER_wrapper.timer
+        logger.info(f"Remain time: {timer.remain_time_duration}")
+
+        if timer.remain_time_duration >= timedelta(hours=DS_RD_SETTING.merge_hours):
+
+            if DS_RD_SETTING.enable_inject_knowledge_at_root:
+
+                if len(trace.hist) == 0:
+                    # set the knowledge base option to True for the first trace
+                    DS_RD_SETTING.enable_knowledge_base = True
+
+                else:
+                    # set the knowledge base option back to False for the other traces
+                    DS_RD_SETTING.enable_knowledge_base = False
+            return self.exp_gen.gen(trace)
+        else:
+            # disable reset in merging stage
+            DS_RD_SETTING.coding_fail_reanalyze_threshold = 100000
+            DS_RD_SETTING.consecutive_errors = 100000
+
+            leaves: list[int] = trace.get_leaves()
+            if len(leaves) < 2:
+                trace.set_current_selection(selection=(-1,))
+                return self.exp_gen.gen(trace)
+            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],)
+                trace.set_current_selection(selection)
+                return self.merge_exp_gen.gen(trace)
@@ -22,6 +22,72 @@ task: |-
  {% if exp_to_merge_fb_desc %}
  {{ exp_to_merge_fb_desc }}
  {% endif %}
+trace: |-
+    {% if exp_and_feedback_list|length <= 1 %}
+    No previous {% if type == "success" %}SOTA{% elif type == "failure" %}failed{% endif %} experiments available.
+    {% else %}
+    {% for exp_and_feedback in exp_and_feedback_list[1:] %}
+    ## Experiment Index: {{ loop.index }}
+    Target Problem: {{ exp_and_feedback[0].hypothesis.problem_desc }}
+    {% if not pipeline %}Chosen Component: {{ exp_and_feedback[0].hypothesis.component }}{% endif %}
+    Proposed Hypothesis: {{ exp_and_feedback[0].hypothesis.hypothesis }}
+    Surpass Previous SOTA: {{ exp_and_feedback[1].decision }}
+    {% if exp_and_feedback[0].result is none %}
+    Experiment Score: Running buggy
+    Experiment Error: {{ exp_and_feedback[1].reason }}
+    {% else %}
+    Experiment Score: {{ exp_and_feedback[0].result.loc["ensemble"].iloc[0] }}
+    Experiment Feedback: {{ exp_and_feedback[1].reason }}
+    {% endif %}
+    {% endfor %}
+    {% endif %}
+hypothesis_gen:
+  system: |-
+    {% include "scenarios.data_science.share:scen.role" %}
+    The user is improving a Kaggle competition implementation iteratively through traces where each new trace is modified from the current SOTA in the trace. If new trace surpasses the current SOTA, it will be the new SOTA. If not, it will be a failed experiment.
+    You will be provided with: 
+      1. A detailed competition scenario description;
+      2. Previous SOTA experiments and feedbacks, which are past SOTA experiments indexed from oldest to newest;
+      3. The current SOTA implementation and feedback, which is the latest SOTA experiments from the previous experiments;
+      4. Extra implementations from another users' experiments;
+    Your task is to:
+      1. **Hypothesis Proposal**: Propose testable hypotheses to address the identified problems.
+      2. **Hypothesis Evaluation**: Evaluate the proposed hypotheses across multiple dimensions.
+
+    # Task 1: Hypothesis Proposal
+    For each identified problem, propose a hypothesis to improve the current SOTA implementation.
+
+    ## Hypothesis Guidelines
+    Here are few guidelines to help you formulate hypotheses:
+    1. Previous Experiments Analysis
+      - For previous SOTA experiments, analyze insights and implicit patterns that can be leveraged to improve the current SOTA implementation.
+      - For failed experiments, think about the persistent problems they facing. If these experiments consistently failed due to time/memory constraints, prioritize changes on efficiency.
+    2. Note on Time/Memory Constraints
+      - If prior experiments failed due to time/memory limitations, assume your new hypothesis will face the same constraints. In this case, prioritize efficiency and **ONLY** response to the problems related to time/memory constraints in your response dictionary.
+      - Besides, do not compromise performance merely for efficiency since the current SOTA implementation do not encounter the constraints. You should think about how to balance the efficiency and performance so that your new hypothesis can be executed successfully and achieve satisfactory performance. 
+
+    # Task 2: Hypothesis Evaluation
+    ## Evaluation Instruction
+    Firstly, you should tag the hypothesis with one of the following components. If the hypothesis is related to multiple components, you should choose the most relevant one.
+    {{ component_desc }}
+    After proposing the hypothesis, your second task is to evaluate the hypothesis from multiple dimensions.
+
+    Secondly, please score the proposed hypothesis from 1 to 10 for each of the following dimensions (where 1 means lowest and 10 means highest):
+    1. Problem-Hypothesis Alignment: How well the hypothesis addresses the identified problem.
+    2. Expected Impact: The estimated improvement after applying the hypothesis to current SOTA implementation.
+    3. Novelty: Degree of innovation compared to previous attempts. If the proposed hypothesis is similar to previous experiments' hypothesis, assign novelty score to one.
+    4. Feasibility: The ease of implementing the proposed hypothesis in the current SOTA implementation.
+    5. Risk-Reward Balance: The exploration-exploitation balance of the proposed hypothesis.
+
+    ## Final Output Format in JSON Schema:
+    {{ hypothesis_output_format }}
+    
+  user: |-
+    # Ertra Experiments and Feedbacks
+    {{ exp_and_feedback_list_desc }}
+
+    # Current SOTA Implementation
+    {{ sota_exp_desc }}

 multi_trace: |-
  {% include "scenarios.data_science.share:scen.role" %}
@@ -52,3 +118,4 @@ multi_trace: |-
  {% endif %}

  {% endfor %}
+ 
@@ -1,7 +1,7 @@
 import json
 import pprint
 from enum import Enum
-from typing import Dict, List, Tuple
+from typing import Dict, List, Optional, Tuple

 import pandas as pd
 from pydantic import BaseModel, Field
@@ -524,6 +524,8 @@ class DSProposalV2ExpGen(ExpGen):
    ) -> Dict:
        problem_formatted_str = ""
        for problem_name, problem_dict in problems.items():
+            if "problem" not in problem_dict:
+                continue
            problem_formatted_str += f"Problem Name: {problem_name}\n"
            problem_formatted_str += f"- Problem Description: {problem_dict['problem']}\n"
            if "idea" in problem_dict:
@@ -557,13 +559,12 @@ class DSProposalV2ExpGen(ExpGen):
        resp_dict = json.loads(response)
        return resp_dict

-    def hypothesis_rank(
+    def compute_top_scores(
        self,
        hypothesis_dict: dict,
-        problem_dict: dict,
-    ) -> Tuple[str, DSHypothesis]:
+    ) -> pd.Series:
        """
-        This function depends on the `identify_problem` function.
+        Compute weighted total scores for each hypothesis and return the top five.
        """
        weights = {
            "alignment_score": 0.2,
@@ -590,8 +591,19 @@ class DSProposalV2ExpGen(ExpGen):

        scores = pd.DataFrame(scores_dict)
        scores_sorted = scores.sum().sort_values(ascending=False)
-        scores_sorted = scores_sorted[:5]  # Select top 5 hypotheses
+        return scores_sorted[:5]

+    def select_hypothesis(
+        self,
+        scores_sorted: pd.Series,
+        hypothesis_dict: dict,
+        problem_dict: dict,
+    ) -> int:
+        """
+        From the top five hypotheses (by weighted score), select one based on additional weighting rules
+        for 'inspired' flag and 'SCENARIO_PROBLEM' label. Returns the chosen hypothesis name and a
+        DSHypothesis instance.
+        """
        # Increase the weight of the hypothesis that is inspired by the idea pool to 3x.
        # Linear decay the weight of the scenario problem from 3x to 0x.
        index_to_pick_pool_list = []
@@ -608,7 +620,21 @@ class DSProposalV2ExpGen(ExpGen):
        reproducible_int = int.from_bytes(bytes.fromhex(md5_hash(scores_sorted.to_string())), byteorder="big") % len(
            index_to_pick_pool_list
        )
-        selected_idx = index_to_pick_pool_list[reproducible_int]
+        return index_to_pick_pool_list[reproducible_int]
+
+    def hypothesis_rank(
+        self, hypothesis_dict: dict, problem_dict: dict, selected_idx: Optional[int] = None
+    ) -> Tuple[str, DSHypothesis]:
+        """
+        Wrapper method that computes the top five hypotheses by weighted scoring and then selects one
+        according to additional weighting rules.
+        """
+        scores_sorted = self.compute_top_scores(hypothesis_dict)
+        if selected_idx is None:
+            selected_idx = self.select_hypothesis(
+                scores_sorted=scores_sorted, hypothesis_dict=hypothesis_dict, problem_dict=problem_dict
+            )
+
        max_score_problem_name = scores_sorted.index[selected_idx]
        problem_dict = problem_dict.get(max_score_problem_name, {})

@@ -737,7 +763,7 @@ class DSProposalV2ExpGen(ExpGen):

        # Step 1: Identify problems
        all_problems = self.identify_problem(
-            current_sub_trace=trace.collect_all_ancestors(),
+            current_sub_trace=trace.get_parent_exps(),
            scenario_desc=scenario_desc,
            sota_exp_desc=sota_exp_desc,
            exp_feedback_list_desc=exp_feedback_list_desc,
@@ -11,15 +11,19 @@ import importlib
 import json
 import re
 import sys
-from multiprocessing import Process, Queue
 from pathlib import Path
 from types import ModuleType
 from typing import Union

+import regex  # type: ignore[import-untyped]
+
 from rdagent.log import rdagent_logger as logger
 from rdagent.oai.llm_conf import LLM_SETTINGS
 from rdagent.utils.agent.tpl import T

+# Default timeout (in seconds) for all regex operations
+REGEX_TIMEOUT = 120.0
+

 def get_module_by_module_path(module_path: Union[str, ModuleType]) -> ModuleType:
    """Load module from path like a/b/c/d.py or a.b.c.d
@@ -68,114 +72,103 @@ def convert2bool(value: Union[str, bool]) -> bool:
        raise ValueError(f"Unknown value type {value} to bool")


-def remove_ansi_codes(s: str) -> str:
+def try_regex_sub(pattern: str, text: str, replace_with: str = "", flags: int = 0) -> str:
    """
-    It is for removing ansi ctrl characters in the string(e.g. colored text)
+    Try to sub a regex pattern against a text string.
    """
-    ansi_escape = re.compile(r"\x1B\[[0-?]*[ -/]*[@-~]")
-    return ansi_escape.sub("", s)
-
-
-def safe_sub(pattern: str, text: str, queue: Queue) -> None:
    try:
-        result = re.sub(pattern, "", text)
-        queue.put(result)
+        text = regex.sub(pattern, replace_with, text, timeout=REGEX_TIMEOUT, flags=flags)
+    except TimeoutError:
+        logger.warning(f"Pattern '{pattern}' timed out after {REGEX_TIMEOUT} seconds; skipping it.")
    except Exception as e:
-        queue.put(e)
+        logger.warning(f"Pattern '{pattern}' raised an error: {e}; skipping it.")
+    return text


-def apply_regex_with_timeout(pattern: str, text: str, timeout: int = 120) -> str:
-    queue: Queue[str | Exception] = Queue()
-    p = Process(target=safe_sub, args=(pattern, text, queue))
-    p.start()
-    p.join(timeout)
-
-    if p.is_alive():
-        p.terminate()
-        p.join()
-        logger.warning(f"Pattern {pattern} timed out after {timeout} seconds, skipping it.")
-        return text
-    else:
-        result = queue.get()
-        if isinstance(result, Exception):
-            logger.warning(f"Pattern {pattern} raised an error: {result}")
-            return text
-        return result
-
-
-def filter_with_time_limit(regex_patterns: Union[str, list[str]], filtered_stdout: str) -> str:
-    if isinstance(regex_patterns, list):
-        for pattern in regex_patterns:
-            filtered_stdout = apply_regex_with_timeout(pattern, filtered_stdout)
-    else:
-        filtered_stdout = re.sub(regex_patterns, "", filtered_stdout)
-    return filtered_stdout
+def filter_with_time_limit(regex_patterns: Union[str, list[str]], text: str) -> str:
+    """
+    Apply one or more regex patterns to filter `text`, using a timeout for each substitution.
+    If `regex_patterns` is a list, they are applied sequentially; if a single string, only that pattern is applied.
+    """
+    if not isinstance(regex_patterns, list):
+        regex_patterns = [regex_patterns]
+    for pattern in regex_patterns:
+        text = try_regex_sub(pattern, text)
+    return text


 def filter_redundant_text(stdout: str) -> str:
    """
-    Filter out progress bars from stdout using regex.
+    Filter out progress bars and other redundant patterns from stdout using regex-based trimming.
    """
    from rdagent.oai.llm_utils import APIBackend  # avoid circular import

-    # Initial progress bar regex pattern
-    progress_bar_re = (
-        r"(\d+/\d+\s+[━]+\s+\d+s?\s+\d+ms/step.*?\u0008+|"
-        r"\d+/\d+\s+[━]+\s+\d+s?\s+\d+ms/step|"
-        r"\d+/\d+\s+[━]+\s+\d+s?\s+\d+ms/step.*|"
-        r"\d+/\d+\s+[━]+.*?\u0008+|"
-        r"\d+/\d+\s+[━]+.*|[ ]*\u0008+|"
-        r"\d+%\|[█▏▎▍▌▋▊▉]+\s+\|\s+\d+/\d+\s+\[\d{2}:\d{2}<\d{2}:\d{2},\s+\d+\.\d+it/s\]|"
-        r"\d+%\|[█]+\|\s+\d+/\d+\s+\[\d{2}:\d{2}<\d{2}:\d{2},\s*\d+\.\d+it/s\])"
-    )
+    # Compile a regex that matches common progress‐bar patterns
+    progress_bar_pattern = r"""(
+        \d+/\d+\s+[━]+\s+\d+s?\s+\d+ms/step.*?\u0008+ |  # e.g. "10/100 ━━━━━━ 3s 50ms/step"
+        \d+/\d+\s+[━]+\s+\d+s?\s+\d+ms/step |            # e.g. "10/100 ━━━━━━ 3s 50ms/step" (no backspaces)
+        \d+/\d+\s+[━]+\s+\d+s?\s+\d+ms/step.* |           # e.g. partial lines
+        \d+/\d+\s+[━]+.*?\u0008+ |                       # e.g. with backspaces
+        \d+/\d+\s+[━]+.* |                                # e.g. partial bars
+        [ ]*\u0008+ |                                     # stray backspaces
+        \d+%\|[█▏▎▍▌▋▊▉]+\s+\|\s+\d+/\d+\s+\[\d{2}:\d{2}<\d{2}:\d{2},\s+\d+\.\d+it/s\] |  # tqdm‐style
+        \d+%\|[█]+\|\s+\d+/\d+\s+\[\d{2}:\d{2}<\d{2}:\d{2},\s*\d+\.\d+it/s\]
+    )"""

-    filtered_stdout = remove_ansi_codes(stdout)
-    filtered_stdout = re.sub(progress_bar_re, "", filtered_stdout)
-    filtered_stdout = re.sub(r"\s*\n\s*", "\n", filtered_stdout)
+    filtered_stdout = try_regex_sub(r"\x1B\[[0-?]*[ -/]*[@-~]", stdout)
+    filtered_stdout = try_regex_sub(progress_bar_pattern, filtered_stdout, flags=regex.VERBOSE)

-    needs_sub = True
-    # Attempt further filtering up to 3 times
+    # Collapse any excessive blank lines/spaces
+    filtered_stdout = try_regex_sub(r"\s*\n\s*", filtered_stdout, replace_with="\n")
+
+    # Iteratively ask the LLM for additional filtering patterns (up to 3 rounds)
    for _ in range(3):
-        filtered_stdout_shortened = filtered_stdout
+        truncated_stdout = filtered_stdout
        system_prompt = T(".prompts:filter_redundant_text.system").r()

-        # Check if all regex patterns have been collected
+        # Try to shrink the stdout so its token count is manageable
        for __ in range(10):
-            user_prompt = T(".prompts:filter_redundant_text.user").r(
-                stdout=filtered_stdout_shortened,
-            )
+            user_prompt = T(".prompts:filter_redundant_text.user").r(stdout=truncated_stdout)
            stdout_token_size = APIBackend().build_messages_and_calculate_token(
                user_prompt=user_prompt,
                system_prompt=system_prompt,
            )
            if stdout_token_size < LLM_SETTINGS.chat_token_limit * 0.1:
-                return filtered_stdout_shortened
+                return truncated_stdout
            elif stdout_token_size > LLM_SETTINGS.chat_token_limit * 0.6:
-                filtered_stdout_shortened = (
-                    filtered_stdout_shortened[: int(LLM_SETTINGS.chat_token_limit * 0.3)]
-                    + filtered_stdout_shortened[-int(LLM_SETTINGS.chat_token_limit * 0.3) :]
-                )
+                head = truncated_stdout[: int(LLM_SETTINGS.chat_token_limit * 0.3)]
+                tail = truncated_stdout[-int(LLM_SETTINGS.chat_token_limit * 0.3) :]
+                truncated_stdout = head + tail
            else:
                break

-        response = json.loads(
-            APIBackend().build_messages_and_create_chat_completion(
-                user_prompt=user_prompt,
-                system_prompt=system_prompt,
-                json_mode=True,
-                json_target_type=dict,
+        try:
+            response = json.loads(
+                APIBackend().build_messages_and_create_chat_completion(
+                    user_prompt=user_prompt,
+                    system_prompt=system_prompt,
+                    json_mode=True,
+                    json_target_type=dict,
+                )
            )
-        )
+        except Exception as e:
+            logger.error(f"LLM filtering request failed: {e}")
+            break
+
        needs_sub = response.get("needs_sub", True)
        regex_patterns = response.get("regex_patterns", [])
-        try:
-            filter_with_time_limit(regex_patterns, filtered_stdout_shortened)

-            if not needs_sub:
-                break
-            filtered_stdout = re.sub(r"\s*\n\s*", "\n", filtered_stdout)
+        try:
+            new_filtered = filter_with_time_limit(regex_patterns, truncated_stdout)
        except Exception as e:
-            logger.error(f"Error in filtering progress bar: due to {e}")
+            logger.error(f"Error applying LLM‐suggested patterns: {e}")
+            break
+
+        if not needs_sub:
+            return new_filtered
+
+        filtered_stdout = try_regex_sub(r"\s*\n\s*", new_filtered, replace_with="\n")
+
    return filtered_stdout