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)