feat: advanced checkpoint selectors (#790)

* rebase selection code

* bug-free run: checkpoint selection and dynamic EDA loading

* add prototypes of various selectors, to imp. and test later

* fix EDA write bug

* imp SOTA-Jump policy

* fix small bug

* allow to set different selector by .env

* add always-win selector

* add init length for AlwaysWinCKPSelector

* add back_jump selector

* auto lint

* add sota_exp_to_submit attribute; change the name of ckp_selector and sota-selector

* fix bug

* auto lint

* working on auto sota selector

* add subtrace counter

* fix bug, remove unuse selector

* add auto sota selector

* auto lint

* fix bug

* fix small logic bug

* add logging

* add inject_diverse feat

* auto lint

* capable to None-select

* feat: add hypothesis_gen config and ExpGen2TraceAndMerge functionality

* refactor: use dynamic import for experiment generator instantiation

* feat: add BestValidSelector for improved SOTA experiment selection

* runnable twin-trace version

* fix logic error of trace-merge

* auto lint

* use import_class to set selector,

* auto-lint

---------

Co-authored-by: Young <afe.young@gmail.com>

rdagent/app/data_science/conf.py

@@ -14,6 +14,9 @@ class DataScienceBasePropSetting(KaggleBasePropSetting):
     scen: str = "rdagent.scenarios.data_science.scen.KaggleScen"
     """Scenario class for data mining model"""
 
+    hypothesis_gen: str = "rdagent.scenarios.data_science.proposal.exp_gen.DSExpGen"
+    """Hypothesis generation class"""
+
     ## Workflow Related
     consecutive_errors: int = 5
 
@@ -47,6 +50,20 @@ class DataScienceBasePropSetting(KaggleBasePropSetting):
     enable_doc_dev: bool = False
     model_dump_check_level: Literal["medium", "high"] = "medium"
 
+    ### selector related
+
+    #### checkpoint selector related
+    # selector_name: str = "latest"
+    selector_name: str = "rdagent.scenarios.data_science.proposal.exp_gen.ckp_select.LatestCKPSelector"
+    """The name of the selector to use"""
+    sota_count_window: int = 5
+    """The number of trials to consider for SOTA count"""
+    sota_count_threshold: int = 1
+    """The threshold for SOTA count"""
+
+    #### SOTA experiment selector related
+    sota_exp_selector_name: str = "rdagent.scenarios.data_science.proposal.exp_gen.sota_exp_select.GlobalSOTASelector"
+    """The name of the SOTA experiment selector to use"""
     ### knowledge base
     enable_knowledge_base: bool = False
     knowledge_base_version: str = "v1"
@@ -65,5 +82,8 @@ class DataScienceBasePropSetting(KaggleBasePropSetting):
     """We'll use f"{DS_RD_SETTING.local_data_path}/{DS_RD_SETTING.eval_sub_dir}/{competition}"
     to find the scriipt to evaluate the submission on test"""
 
+    ### inject diverse
+    enable_inject_diverse: bool = False
+
 
 DS_RD_SETTING = DataScienceBasePropSetting()

rdagent/app/data_science/loop.py

@@ -32,10 +32,31 @@ from rdagent.scenarios.data_science.dev.feedback import DSExperiment2Feedback
 from rdagent.scenarios.data_science.dev.runner import DSCoSTEERRunner
 from rdagent.scenarios.data_science.experiment.experiment import DSExperiment
 from rdagent.scenarios.data_science.proposal.exp_gen import DSExpGen, DSTrace
+from rdagent.scenarios.data_science.proposal.exp_gen.ckp_select import (
+    BackJumpCKPSelector,
+    LatestCKPSelector,
+    SOTAJumpCKPSelector,
+)
 from rdagent.scenarios.data_science.proposal.exp_gen.idea_pool import DSKnowledgeBase
-from rdagent.scenarios.data_science.proposal.exp_gen.select import LatestCKPSelector
+from rdagent.scenarios.data_science.proposal.exp_gen.sota_exp_select import (
+    AutoSOTAexpSelector,
+    BestValidSelector,
+    GlobalSOTASelector,
+)
 from rdagent.scenarios.kaggle.kaggle_crawler import download_data
 
+CKP_SELECTOR_NAME_MAP = {
+    "latest": LatestCKPSelector,
+    "sota_jump": SOTAJumpCKPSelector,
+    "back_jump": BackJumpCKPSelector,
+}
+
+SOTA_EXP_SELECTOR_NAME_MAP = {
+    "global_sota": GlobalSOTASelector,
+    "auto_sota": AutoSOTAexpSelector,
+    "best_valid_sota": BestValidSelector,
+}
+
 
 class DataScienceRDLoop(RDLoop):
     skip_loop_error = (CoderError, RunnerError)
@@ -49,8 +70,15 @@ class DataScienceRDLoop(RDLoop):
 
         # 2) task generation from a complete solution
         # self.exp_gen: ExpGen = import_class(PROP_SETTING.exp_gen)(scen)
-        self.ckp_selector = LatestCKPSelector()
-        self.exp_gen = DSExpGen(scen)
+
+        # self.ckp_selector = CKP_SELECTOR_NAME_MAP[DS_RD_SETTING.selector_name]()
+        # self.sota_exp_selector = SOTA_EXP_SELECTOR_NAME_MAP[DS_RD_SETTING.sota_exp_selector_name]()
+        self.ckp_selector = import_class(PROP_SETTING.selector_name)()
+        self.sota_exp_selector = import_class(PROP_SETTING.sota_exp_selector_name)()
+
+        self.exp_gen = import_class(PROP_SETTING.hypothesis_gen)(scen)
+
+        # coders
         self.data_loader_coder = DataLoaderCoSTEER(scen)
         self.feature_coder = FeatureCoSTEER(scen)
         self.model_coder = ModelCoSTEER(scen)
@@ -76,6 +104,12 @@ class DataScienceRDLoop(RDLoop):
         super(RDLoop, self).__init__()
 
     def direct_exp_gen(self, prev_out: dict[str, Any]):
+
+        # set the SOTA experiment to submit
+        sota_exp_to_submit = self.sota_exp_selector.get_sota_exp_to_submit(self.trace)
+        self.trace.set_sota_exp_to_submit(sota_exp_to_submit)
+
+        # set the checkpoint to start from
         selection = self.ckp_selector.get_selection(self.trace)
         exp = self.exp_gen.gen(self.trace, selection)
         logger.log_object(exp)

rdagent/components/coder/data_science/pipeline/eval.py

@@ -127,6 +127,13 @@ class PipelineCoSTEEREvaluator(CoSTEEREvaluator):
 
         eda_output = implementation.file_dict.get("EDA.md", None)
 
+        eda_output = implementation.file_dict.get("EDA.md", None)
+
+        if not isinstance(implementation, FBWorkspace):
+            eda_output = None
+        else:
+            eda_output = implementation.file_dict.get("EDA.md", None)
+
         system_prompt = T(".prompts:pipeline_eval.system").r(
             scenario=self.scen.get_scenario_all_desc(eda_output=eda_output),
             task_desc=target_task.get_task_information(),

rdagent/core/proposal.py

@@ -149,11 +149,22 @@ class CheckpointSelector:
         - `(idx, )` represents starting from the `idx`-th trial in the trace.
         - `None` represents starting from scratch (start a new trace)
 
-
         - More advanced selection strategies in `select.py`
         """
 
 
+class SOTAexpSelector:
+    """
+    Select the SOTA experiment from the trace to submit
+    """
+
+    @abstractmethod
+    def get_sota_exp_to_submit(self, trace: Trace) -> Experiment | None:
+        """
+        Select the SOTA experiment from the trace to submit
+        """
+
+
 class ExpGen(ABC):
 
     def __init__(self, scen: Scenario) -> None:

rdagent/scenarios/data_science/dev/feedback.py

@@ -31,9 +31,11 @@ class DSExperiment2Feedback(Experiment2Feedback):
             exp=sota_exp, heading="SOTA of previous exploration of the scenario"
         )
 
+        last_exp = trace.last_exp()
+
         # Get feedback description using shared template
         feedback_desc = T("scenarios.data_science.share:describe.feedback").r(
-            exp_and_feedback=(trace.hist[-1] if trace.hist else None), heading="Previous Trial Feedback"
+            exp_and_feedback=trace.hist[-1] if trace.hist else None, heading="Previous Trial Feedback"
         )
 
         # TODO:

rdagent/scenarios/data_science/proposal/exp_gen/base.py

@@ -61,10 +61,17 @@ class DSTrace(Trace[DataScienceScen, KnowledgeBase]):
 
         self.knowledge_base = knowledge_base
 
+        self.sub_trace_count: int = 0
+
         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")
 
+    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
 
@@ -127,15 +134,22 @@ class DSTrace(Trace[DataScienceScen, KnowledgeBase]):
         list[tuple[DSExperiment, ExperimentFeedback]]
             The search list.
         """
+        if search_type == "all":
+            return self.hist
 
-        if selection is None:
-            selection = self.get_current_selection()
+        elif search_type == "ancestors":
 
-        if selection is None:
-            # selection is None, which means we switch to a new trace, which is not implemented yet
-            return []
+            if selection is None:
+                selection = self.get_current_selection()
 
-        return self.collect_all_ancestors(selection) if search_type == "ancestors" else self.hist
+            if len(selection) == 0:
+                # selection is (), which means we switch to a new trace
+                return []
+
+            return self.collect_all_ancestors(selection)
+
+        else:
+            raise ValueError(f"Invalid search type: {search_type}")
 
     def collect_all_ancestors(
         self,

rdagent/scenarios/data_science/proposal/exp_gen/ckp_select.py

@@ -0,0 +1,149 @@
+import random
+
+from rdagent.app.data_science.conf import DS_RD_SETTING
+from rdagent.core.proposal import CheckpointSelector, Trace
+from rdagent.log import rdagent_logger as logger
+
+# # TODO: more advanced selector
+# # TODO/Discussion: load selector function here or define selector class in `proposal.py`?
+
+
+class LatestCKPSelector(CheckpointSelector):
+    """
+    -`(-1, )` represents starting from the latest trial in the trace
+    """
+
+    def __init__(
+        self,
+    ):
+        logger.info(f"Using latest selector by default")
+
+    def get_selection(self, trace: Trace) -> tuple[int, ...]:
+
+        return (-1,)
+
+
+class SOTAJumpCKPSelector(CheckpointSelector):
+    """
+    SOTA jump policy:
+    if the cumulative SOTA in a window is below a threshold, jump to a new trial
+    otherwise, continue the current latest trial
+    """
+
+    def __init__(
+        self,
+    ) -> None:
+        self.SOTA_COUNT_WINDOW = DS_RD_SETTING.sota_count_window
+        self.SOTA_COUNT_THRESHOLD = DS_RD_SETTING.sota_count_threshold
+
+        logger.info(
+            f"Using SOTA-jump selector with window {self.SOTA_COUNT_WINDOW} and threshold {self.SOTA_COUNT_THRESHOLD}"
+        )
+
+    def get_selection(self, trace: Trace) -> tuple[int, ...]:
+
+        current_trace = trace.retrieve_search_list(search_type="ancestors")
+        if len(trace.hist) > 0 and len(current_trace) > self.SOTA_COUNT_WINDOW:
+            all_exp_list = trace.experiment_and_feedback_list_after_init(return_type="all", search_type="ancestors")
+            # sota_exp_list = trace.experiment_and_feedback_list_after_init(return_type="sota", search_type="ancestors")
+            exp_list_in_window = all_exp_list[-self.SOTA_COUNT_WINDOW :]
+
+            # compute the cumulative SOTA ratio in the window
+            sota_count = 0
+            for exp, fb in exp_list_in_window:
+                if fb.decision:
+                    sota_count += 1
+            if sota_count < self.SOTA_COUNT_THRESHOLD:
+                trace.sub_trace_count += 1
+                logger.info(
+                    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 ()
+            else:
+                logger.info(
+                    f"SOTA count {sota_count} is above threshold {self.SOTA_COUNT_THRESHOLD}, continue the current latest trial"
+                )
+                logger.info(f"current sub-trace count: {trace.sub_trace_count}")
+                return (-1,)
+
+        else:
+            logger.info(f"Not enough history to make a decision, continue the current latest trial")
+            return (-1,)
+
+
+class BackJumpCKPSelector(CheckpointSelector):
+    """
+    back-jump policy:
+    if the cumulative SOTA in a window is below a threshold,
+    with 50% probability, reboot a new sub-trace
+    with 50% probability, jump back to the "last second" SOTA trial (we assume the lastest SOTA trial is not good enough selection)
+    """
+
+    def __init__(
+        self,
+    ) -> None:
+        self.SOTA_COUNT_WINDOW = DS_RD_SETTING.sota_count_window
+        self.SOTA_COUNT_THRESHOLD = DS_RD_SETTING.sota_count_threshold
+
+        logger.info(
+            f"Using back-jump selector with window {self.SOTA_COUNT_WINDOW} and threshold {self.SOTA_COUNT_THRESHOLD}"
+        )
+
+    def get_selection(self, trace: Trace) -> tuple[int, ...]:
+        current_trace = trace.retrieve_search_list(search_type="ancestors")
+
+        if len(trace.hist) > 0 and len(current_trace) > self.SOTA_COUNT_WINDOW:
+
+            all_exp_list = trace.experiment_and_feedback_list_after_init(return_type="all", search_type="ancestors")
+            # sota_exp_list = trace.experiment_and_feedback_list_after_init(return_type="sota", search_type="ancestors")
+            exp_list_in_window = all_exp_list[-self.SOTA_COUNT_WINDOW :]
+
+            # compute the cumulative SOTA ratio in the window
+            sota_count = 0
+            for exp, fb in exp_list_in_window:
+                if fb.decision:
+                    sota_count += 1
+
+            if sota_count < self.SOTA_COUNT_THRESHOLD:
+
+                random_choice = random.random()
+                if random_choice < 0.5:
+                    trace.sub_trace_count += 1
+                    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
+                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)"
+                    )
+                    sota_exp_list = trace.experiment_and_feedback_list_after_init(return_type="sota", search_type="all")
+                    if len(sota_exp_list) > 1:
+                        last_second_sota_idx = trace.hist.index(sota_exp_list[-2])
+                        logger.info(
+                            f"jump back to the last second SOTA in hist (may not in current sub-trace), index: {last_second_sota_idx}"
+                        )
+                        logger.info(f"current sub-trace count: {trace.sub_trace_count}")
+                        return (last_second_sota_idx,)
+                    else:
+                        trace.sub_trace_count += 1
+                        logger.info(
+                            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
+
+            else:
+                logger.info(
+                    f"SOTA count {sota_count} is above threshold {self.SOTA_COUNT_THRESHOLD}, continue the current latest trial"
+                )
+                logger.info(f"current sub-trace count: {trace.sub_trace_count}")
+                return (-1,)
+        else:
+            logger.info(f"Not enough history to make a decision, continue the current latest trial")
+            logger.info(f"current sub-trace count: {trace.sub_trace_count}")
+            return (-1,)
+
+
+# TODO: implement these selectors and more

rdagent/scenarios/data_science/proposal/exp_gen/merge.py

@@ -1,9 +1,14 @@
 """Merge the version in different traces"""
 
+from datetime import timedelta
+
 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.scenarios.data_science.experiment.experiment import DSExperiment
+from rdagent.scenarios.data_science.proposal.exp_gen import DSExpGen
 from rdagent.scenarios.data_science.proposal.exp_gen.base import DSHypothesis, DSTrace
 from rdagent.utils.agent.tpl import T
 
@@ -75,3 +80,34 @@ class MergeExpGen(ExpGen):
         if sota_exp_fb is not None:
             exp.experiment_workspace.inject_code_from_file_dict(sota_exp_fb[0].experiment_workspace)
         return exp
+
+
+class ExpGen2TraceAndMerge(ExpGen):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.merge_exp_gen = MergeExpGen(self.scen)
+        self.exp_gen = DSExpGen(self.scen)
+
+    def gen(self, trace: DSTrace, selection: tuple[int, ...] = (-1,)) -> DSExperiment:
+        timer: RDAgentTimer = RD_Agent_TIMER_wrapper.timer
+        logger.info(f"Remain time: {timer.remain_time_duration}")
+
+        if timer.remain_time_duration >= timedelta(hours=2):
+            leaves: list[int] = trace.get_leaves()
+            if len(leaves) < 2:
+                selection = tuple()  # create new trace
+            else:
+                selection = (
+                    leaves[0],
+                )  # continue the first trace. This will result in the interleaving of two traces expansion.
+            return self.exp_gen.gen(trace, selection)
+        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:
+                return self.exp_gen.gen(trace, selection)
+            else:
+                return self.merge_exp_gen.gen(trace, selection)

rdagent/scenarios/data_science/proposal/exp_gen/prompts_selector.yaml

@@ -0,0 +1,24 @@
+auto_sota_selector:
+  system: |-
+    
+    You are a data scientist and a top Kaggle competitor. The user is working on improving a solution for a Kaggle competition. The user has already conducted a series of successful experiments (SOAT trails during the exploration) and collected feedbacks.
+    
+    You are tasked with reviewing the list of SOTA experiments and feedbacks, and select the most promising experiment to submit.
+
+    Please be objective and data-driven in your analysis, and provide a explanation for your selection. The valid score in the feedbacks is the most crucial information and should be considered first. The risk on overfitting should be considered as well.
+
+    # The scenario and the description of the competition are as follows:
+    {{ scenario }}
+
+    # Your response should be short and concise, strictly adhere to the following JSON format:
+    {
+      "selected_SOTA_idx": [Experiment No.](positive integer),
+      "explanation": "A brief explanation text for your selection."
+    }
+
+  user: |-
+    # SOTA Experiments and Feedback
+    {{ historical_sota_exp_with_desc_and_scores }}
+
+
+

rdagent/scenarios/data_science/proposal/exp_gen/prompts_v2.yaml

@@ -44,6 +44,15 @@ feedback_problem:
       4. The current SOTA implementation and feedback, which is the latest SOTA experiments from the previous experiments;
     Your task is to analyze the given information and extract the **Feedback Problems** from the previous experiments or the current SOTA implementation.
 
+    
+    {% if inject_diverse %}
+    ### Focus on Diversity!!
+    Diversity is very critical in the analysis of scenario problems. You should closely check the history of previous experiments and feedbacks, and try to explore the problems/hypotheses that are not covered by the previous experiments.
+    1. Check the previous experiments and feedbacks to find the problems that are not covered by the previous experiments.
+    2. Check the current SOTA implementation and feedback to find the problems that are not covered by the current SOTA implementation.
+    3. Do not do incremental exploration on the previous problems.
+    {% endif %}
+
     ## Feedback Problems
     ### Definition
     Feedback problems are specific and fine-grained technical, or methodological issues within the previous experiments or the current SOTA implementation.
@@ -143,6 +152,16 @@ hypothesis_gen:
     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.
 
+    {% if inject_diverse %}
+    # Focus on Diversity!!
+    Diversity is very critical in the analysis of scenario problems. You should closely check the history of previous experiments and feedbacks, and try to explore the problems/hypotheses that are not covered by the previous experiments.
+    1. Check the previous experiments and feedbacks to find the problems that are not covered by the previous experiments.
+    2. Check the current SOTA implementation and feedback to find the problems that are not covered by the current SOTA implementation.
+    3. Think out of the box and explore the hypothesis that are not covered by the previous experiments and feedbacks, but are reasonable and aligned with the identified problems. 
+    4. Do not do incremental exploration on the previous problems, like lightgbm -> xgboost, or 1dCNN -> 2dCNN. Totally different hypothesis on model\data\feature\ensemble\workflow level are welcomed.
+    {% endif %}
+
+
     ## Final Output Format in JSON Schema:
     {{ hypothesis_output_format }}
     
@@ -311,4 +330,6 @@ output_format:
     {
       "problem name 1 (should be exactly same as the problem name provided)": 1, # The index which is same to the idea index provided in the input and must be integer.
       "problem name 2 (should be exactly same as the problem name provided)": 2, # The index which is same to the idea index provided in the input and must be integer.
-    }
+    }
+
+

rdagent/scenarios/data_science/proposal/exp_gen/proposal.py

@@ -242,10 +242,13 @@ class DSProposalV2ExpGen(ExpGen):
         )
         return json.loads(response)
 
-    def identify_feedback_problem(self, scenario_desc: str, exp_feedback_list_desc: str, sota_exp_desc: str) -> Dict:
+    def identify_feedback_problem(
+        self, scenario_desc: str, exp_feedback_list_desc: str, sota_exp_desc: str, inject_diverse: bool = False
+    ) -> Dict:
         sys_prompt = T(".prompts_v2:feedback_problem.system").r(
             problem_spec=T(".prompts_v2:specification.problem").r(),
             problem_output_format=T(".prompts_v2:output_format.problem").r(),
+            inject_diverse=inject_diverse,
         )
         user_prompt = T(".prompts_v2:feedback_problem.user").r(
             scenario_desc=scenario_desc,
@@ -270,6 +273,7 @@ class DSProposalV2ExpGen(ExpGen):
         problems: dict,
         pipeline: bool,
         enable_idea_pool: bool,
+        inject_diverse: bool = False,
     ) -> Dict:
         problem_formatted_str = ""
         for problem_name, problem_dict in problems.items():
@@ -288,6 +292,7 @@ class DSProposalV2ExpGen(ExpGen):
             ),
             pipeline=pipeline,
             enable_idea_pool=enable_idea_pool,
+            inject_diverse=inject_diverse,
         )
         user_prompt = T(".prompts_v2:hypothesis_gen.user").r(
             scenario_desc=scenario_desc,
@@ -435,6 +440,7 @@ class DSProposalV2ExpGen(ExpGen):
         return exp
 
     def gen(self, trace: DSTrace, pipeline: bool = False) -> DSExperiment:
+
         if pipeline:
             component_desc = T("scenarios.data_science.share:component_description_in_pipeline").r()
         else:
@@ -467,6 +473,26 @@ class DSProposalV2ExpGen(ExpGen):
             pipeline=pipeline,
         )
 
+        if DS_RD_SETTING.enable_inject_diverse and len(trace.hist) > 0:
+            if len(trace.current_selection) == 0:
+                # start a new sub-trace, and inject diverse problems.
+                inject_diverse = True
+                logger.info("Start a new sub-trace, and inject diverse problems.")
+            else:
+                inject_diverse = False
+        else:
+            inject_diverse = False
+
+        if DS_RD_SETTING.enable_inject_diverse and len(trace.hist) > 0:
+            if len(trace.current_selection) == 0:
+                # start a new sub-trace, and inject diverse problems.
+                inject_diverse = True
+                logger.info("Start a new sub-trace, and inject diverse problems.")
+            else:
+                inject_diverse = False
+        else:
+            inject_diverse = False
+
         # Step 1: Identify problems
         all_problems = {}
         if len(trace.hist) >= 3:
@@ -474,6 +500,7 @@ class DSProposalV2ExpGen(ExpGen):
                 scenario_desc=scenario_desc,
                 exp_feedback_list_desc=exp_feedback_list_desc,
                 sota_exp_desc=sota_exp_desc,
+                inject_diverse=inject_diverse,
             )
             for problem_name in fb_problems:
                 fb_problems[problem_name]["label"] = "FEEDBACK_PROBLEM"
@@ -507,6 +534,7 @@ class DSProposalV2ExpGen(ExpGen):
             problems=all_problems,
             pipeline=pipeline,
             enable_idea_pool=DS_RD_SETTING.enable_knowledge_base,
+            inject_diverse=inject_diverse,
         )
         if not pipeline:
             sota_exp_model_file_count = len(

rdagent/scenarios/data_science/proposal/exp_gen/select.py

@@ -1,111 +0,0 @@
-from rdagent.core.proposal import CheckpointSelector, Trace
-
-# # TODO: more advanced selector
-# # TODO/Discussion: load selector function here or define selector class in `proposal.py`?
-
-
-class LatestCKPSelector(CheckpointSelector):
-    """
-    -`(-1, )` represents starting from the latest trial in the trace
-    """
-
-    def get_selection(self, trace: Trace) -> tuple[int, ...]:
-
-        return (-1,)
-
-
-class SOTAJumpCKPSelector(CheckpointSelector):
-    """
-    SOTA jump policy:
-    if the cumulative SOTA in a window is below a threshold, jump to a new trial
-    otherwise, continue the current latest trial
-    """
-
-    def __init__(self) -> None:
-
-        self.SOTA_COUNT_WINDOW = 5
-        self.SOTA_COUNT_THRESHOLD = 1  # start to compute cumulative SOTA ratio after 5 trials
-
-    def get_selection(self, trace: Trace) -> tuple[int, ...] | None:
-
-        current_trace = trace.retrieve_search_list(search_type="ancestors")
-        if len(trace.hist) > 0 and len(current_trace) > self.SOTA_COUNT_WINDOW:
-            all_exp_list = trace.experiment_and_feedback_list_after_init(return_type="all", search_type="ancestors")
-            # sota_exp_list = trace.experiment_and_feedback_list_after_init(return_type="sota", search_type="ancestors")
-            exp_list_in_window = all_exp_list[-self.SOTA_COUNT_WINDOW :]
-
-            # compute the cumulative SOTA ratio in the window
-            sota_count = 0
-            for exp, fb in exp_list_in_window:
-                if fb.decision:
-                    sota_count += 1
-            if sota_count < self.SOTA_COUNT_THRESHOLD:
-                return None
-            else:
-                return (-1,)
-
-        else:
-            return (-1,)
-
-
-# TODO: implement these selectors and more
-
-
-class GlobalGreedyCKPSelector(CheckpointSelector):
-    """
-    global greedy selector: select the trial with best performance globally (in trace.hist)
-    consistent with the greedy strategy in AIDE
-    not implemented yet
-    """
-
-    def get_selection(self, trace: Trace) -> tuple[int, ...]:
-
-        return (-1,)
-
-
-class LocalGreedyCKPSelector(CheckpointSelector):
-    """
-    local greedy selector: select the trial with best performance locally (in trace.ancestors)
-    not implemented yet
-    """
-
-    def get_selection(self, trace: Trace) -> tuple[int, ...]:
-
-        return (-1,)
-
-
-class BugBufferCKPSelector(CheckpointSelector):
-    """
-    bug buffer selector: with limit-size bug buffer size, start a new trace if buffer exceeds.
-    not implemented yet
-    """
-
-    def __init__(self) -> None:
-        self.bug_count = 0
-        self.BUG_BUFFER_SIZE = 10
-
-    def get_selection(self, trace: Trace) -> tuple[int, ...]:
-
-        if self.bug_count < self.BUG_BUFFER_SIZE:
-            return (-1,)
-
-        else:
-            return None
-
-
-class RandomCKPSelector(CheckpointSelector):
-    def get_selection(self, trace: Trace) -> tuple[int, ...]:
-        """
-        random selector: select the trial randomly
-        not implemented yet
-        """
-        return (-1,)
-
-
-class BuggyCKPSelector(CheckpointSelector):
-    def get_selection(self, trace: Trace) -> tuple[int, ...]:
-        """
-        buggy selector: select the most recent trial with buggy performance
-        not implemented yet
-        """
-        return (-1,)

rdagent/scenarios/data_science/proposal/exp_gen/sota_exp_select.py

@@ -0,0 +1,127 @@
+import json
+import random
+from typing import Dict, Tuple
+
+import numpy as np
+import pandas as pd
+
+from rdagent.app.data_science.conf import DS_RD_SETTING
+from rdagent.core.proposal import ExperimentFeedback, SOTAexpSelector, Trace
+from rdagent.log import rdagent_logger as logger
+from rdagent.oai.llm_utils import APIBackend, md5_hash
+from rdagent.scenarios.data_science.experiment.experiment import DSExperiment
+from rdagent.scenarios.data_science.proposal.exp_gen.base import DSHypothesis, DSTrace
+from rdagent.utils.agent.tpl import T
+from rdagent.utils.workflow import wait_retry
+
+
+class GlobalSOTASelector(SOTAexpSelector):
+    """
+    return the latest SOTA experiment from the trace to submit
+    """
+
+    def __init__(
+        self,
+    ):
+        print(f"Using global SOTA policy by default")
+
+    def get_sota_exp_to_submit(self, trace: Trace) -> DSExperiment | None:
+
+        return trace.sota_experiment(search_type="all")
+
+
+class AutoSOTAexpSelector(SOTAexpSelector):
+    """
+    retrieve a list of SOTA experiments from the trace, then call the LLM to select the best one
+    """
+
+    def __init__(
+        self,
+    ):
+        print(f"Using auto SOTA policy")
+
+    @wait_retry(retry_n=5)
+    def get_sota_exp_to_submit(self, trace: Trace) -> DSExperiment | None:
+        # retrieve all SOTA experiments from the trace
+
+        sota_exp_fb_list = trace.experiment_and_feedback_list_after_init(return_type="sota", search_type="all")
+
+        if len(sota_exp_fb_list) == 0:
+            logger.info("Auto SOTA selector: No SOTA in trace yet")
+            return None
+
+        elif len(sota_exp_fb_list) == 1:
+            sota_idx_in_trace = trace.hist.index(sota_exp_fb_list[0])
+            logger.info(
+                f"Auto SOTA selector: Only one SOTA in trace, using it, which is the No. {sota_idx_in_trace + 1} in the trace"
+            )
+            return sota_exp_fb_list[0][0]
+
+        else:
+            logger.info("Auto SOTA selector: Multiple SOTA in trace, calling LLM to select the best one")
+
+            SOAT_exp_with_desc_and_scores = "Historical SOTA experiments:\n\n"
+
+            for i, (exp, ef) in enumerate(sota_exp_fb_list):
+                if exp:
+                    current_final_score = pd.DataFrame(exp.result).loc["ensemble"].iloc[0]
+                    desc = T("scenarios.data_science.share:describe.exp").r(
+                        exp=exp, heading="SOTA of previous exploration of the scenario"
+                    )
+                    SOAT_exp_with_desc_and_scores += f"""SOTA experiment No. {i+1}:
+                        Description: {desc}
+                        Final score: {current_final_score}\n\n"""
+
+            system_prompt = T(".prompts_selector:auto_sota_selector.system").r(
+                scenario=trace.scen.get_scenario_all_desc()
+            )
+
+            user_prompt = T(".prompts_selector:auto_sota_selector.user").r(
+                historical_sota_exp_with_desc_and_scores=SOAT_exp_with_desc_and_scores,
+            )
+
+            response = APIBackend().build_messages_and_create_chat_completion(
+                user_prompt=user_prompt,
+                system_prompt=system_prompt,
+                json_mode=True,
+                json_target_type=Dict[str, str | int],
+            )
+
+            response_dict = json.loads(response)
+
+            sota_submit_idx = response_dict.get("selected_SOTA_idx", None)
+
+            if sota_submit_idx is not None:
+                sota_submit = sota_exp_fb_list[int(sota_submit_idx) - 1]
+                sota_idx_in_trace = trace.hist.index(sota_submit)
+                logger.info(
+                    f"Auto SOTA selector: selected SOTA experiment No. {sota_submit_idx} to submit, which is the No. {sota_idx_in_trace + 1} in the trace"
+                )
+                return sota_submit[0]
+            else:
+                # no SOTA experiment to submit, using the latest SOTA experiment
+                logger.info("Auto SOTA selector: No SOTA experiment to submit, using the latest SOTA experiment")
+                return sota_exp_fb_list[-1][0]
+
+
+class BestValidSelector(SOTAexpSelector):
+    def get_sota_exp_to_submit(self, trace: Trace) -> DSExperiment | None:
+        sota_exp_fb_list = trace.experiment_and_feedback_list_after_init(return_type="all", search_type="all")
+        direction_sign = 1 if trace.scen.metric_direction else -1
+
+        def get_sort_key(exp_fb: tuple[DSExperiment, ExperimentFeedback]) -> tuple[bool, float]:
+            score = -np.inf
+            result: pd.DataFrame | None = exp_fb[0].result
+            if result is not None:
+                score = result.loc["ensemble"].iloc[0]
+            return (exp_fb[1].decision, direction_sign * score)
+
+        if len(sota_exp_fb_list) == 0:
+            logger.info("Best Valid SOTA selector: No SOTA in trace yet")
+            return None
+        else:
+            sota_exp_fb_list = sorted(sota_exp_fb_list, key=get_sort_key, reverse=True)
+            return sota_exp_fb_list[0][0]
+
+
+# TODO: more advanced sota exp selector (e.g. LLM-based, merge exp with multiple sub-trace)