feat: enable to set different version of idea-proposal for multi traces (#895)

* fix the logic of kb-inject, allow different verion

* set more flexiable proposal-version change for multi-tarce

* auto-lint

* fix the divede-zero-bug in a trival way

* keep the dump imp. first, update in next version

* use get_sub_trace_count() to get trace_num_count

* fix the conern case bug of divide-zero

* update corner case

* fix the bug

* auto-lint

* fis the bug

* fix the logic bug in max_sota_filter

* fix bug of old version of self.exp_gen.gen

* update the reset_exp_gen_version

* use get_parent_exps to replace all  collect_all_ancestors

* auto lint

* fix the bug of reset_exp_gen_version

* fix bug: update V3's old hypothesis_rank

* trival patch on gap of V3 & V2

* make dump patch to unify proposal_V3's dentify_problems

* auto-lint

* fix the bug of sub_trace_count

Author: xuangu-fang

rdagent/app/data_science/conf.py

@@ -91,9 +91,13 @@ class DataScienceBasePropSetting(KaggleBasePropSetting):
     # inject diverse when start a new sub-trace
     enable_inject_diverse: bool = False
 
-    # inject diverse at the root of the trace
+    # inject knowledge at the root of the trace
     enable_inject_knowledge_at_root: bool = False
 
+    # enable different version of DSExpGen for multi-trace
+    enable_multi_version_exp_gen: bool = False
+    exp_gen_version_list: str = "v3,v2"
+
     #### multi-trace: time for final multi-trace merge
     merge_hours: int = 2
     """The time for merge"""

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

@@ -304,21 +304,40 @@ def __init__(self, *args, **kwargs):
         )
         self.flag_start_merge = False
 
-    def gen(self, trace: DSTrace) -> DSExperiment:
+    def reset_exp_gen_version(self, version: str = "v2"):
+        DS_RD_SETTING.proposal_version = version
+        logger.info(f"ExpGen2TraceAndMergeV2: Resetting proposal version to {version}")
+        self.exp_gen = DataScienceRDLoop._get_exp_gen(
+            f"rdagent.scenarios.data_science.proposal.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=DS_RD_SETTING.merge_hours):
 
             if DS_RD_SETTING.enable_inject_knowledge_at_root:
+                if DS_RD_SETTING.knowledge_base_path is not None and DS_RD_SETTING.idea_pool_json_path is not None:
+                    if len(trace.hist) == 0:
+                        # set the knowledge base option to True for the first trace
+                        DS_RD_SETTING.enable_knowledge_base = True
+
+            if DS_RD_SETTING.enable_multi_version_exp_gen:
+                exp_gen_version_list = DS_RD_SETTING.exp_gen_version_list.split(",")
+                for version in exp_gen_version_list:
+                    assert version in ["v3", "v2", "v1"]
 
                 if len(trace.hist) == 0:
-                    # set the knowledge base option to True for the first trace
-                    DS_RD_SETTING.enable_knowledge_base = True
+                    # set the proposal version for the first sub-trace
+                    self.reset_exp_gen_version(version=exp_gen_version_list[0])
+                elif len(trace.get_current_selection()) == 0 and trace.sub_trace_count > 0:
+                    # reset the proposal version at the start of other sub-trace
+                    if trace.sub_trace_count - 1 < len(exp_gen_version_list):
+                        self.reset_exp_gen_version(version=exp_gen_version_list[trace.sub_trace_count - 1])
+                    else:
+                        self.reset_exp_gen_version(version=exp_gen_version_list[-1])
 
-                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:

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

@@ -16,6 +16,12 @@ auto_sota_selector:
       "explanation": "A brief explanation text for your selection."
     }
 
+    If you cannot make a selection, like no SOTA experiments and feedbacks, or the gap is too small, return 
+      {
+        "selected_SOTA_idx": None,
+        "explanation": "No SOTA experiments and feedbacks"
+      }
+
   user: |-
     # SOTA Experiments and Feedback
     {{ historical_sota_exp_with_desc_and_scores }}

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

@@ -1030,6 +1030,46 @@ def get_all_hypotheses(self, problem_dict: dict, hypothesis_dict: dict) -> list[
             )
         return result
 
+    # FIXME: remove this, dump solution, should be merged into identify_problem in V2
+    def identify_problems_v3(
+        self, trace: DSTrace, scenario_desc: str, sota_exp_desc: str, exp_feedback_list_desc: str
+    ) -> Dict:
+        sub_trace = trace.get_parent_exps()
+        trace_length = len(trace.hist)
+        all_problems = {}
+
+        # 阶段一：探索期（主要场景问题）
+        if trace_length <= 3:
+            scen_problems = self.identify_scenario_problem(scenario_desc, sota_exp_desc)
+            for problem_name in scen_problems:
+                scen_problems[problem_name]["label"] = "SCENARIO_PROBLEM"
+                all_problems[problem_name] = scen_problems[problem_name]
+            self.scen_prob_multiplier = 3
+
+        # 阶段二：混合期（两种问题都考虑）
+        elif trace_length <= 6:
+            # 优先场景问题，但也考虑反馈
+            scen_problems = self.identify_scenario_problem(scenario_desc, sota_exp_desc)
+            for problem_name in scen_problems:
+                scen_problems[problem_name]["label"] = "SCENARIO_PROBLEM"
+                all_problems[problem_name] = scen_problems[problem_name]
+
+            fb_problems = self.identify_feedback_problem(scenario_desc, exp_feedback_list_desc, sota_exp_desc)
+            for problem_name in fb_problems:
+                fb_problems[problem_name]["label"] = "FEEDBACK_PROBLEM"
+                all_problems[problem_name] = fb_problems[problem_name]
+            self.scen_prob_multiplier = 2
+
+        # 阶段三：优化期（主要反馈问题）
+        else:
+            fb_problems = self.identify_feedback_problem(scenario_desc, exp_feedback_list_desc, sota_exp_desc)
+            for problem_name in fb_problems:
+                fb_problems[problem_name]["label"] = "FEEDBACK_PROBLEM"
+                all_problems[problem_name] = fb_problems[problem_name]
+            self.scen_prob_multiplier = 1
+
+        return all_problems
+
     def gen(self, trace: DSTrace) -> DSExperiment:
         pipeline = DS_RD_SETTING.coder_on_whole_pipeline
         if not pipeline and (draft_exp := draft_exp_in_decomposition(self.scen, trace)):
@@ -1067,26 +1107,43 @@ def gen(self, trace: DSTrace) -> DSExperiment:
             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
         # Step 1: Identify problems
         all_problems = {}
-        if len(trace.hist) >= 3:
-            fb_problems = self.identify_feedback_problem(
-                scenario_desc=scenario_desc,
-                exp_feedback_list_desc=exp_feedback_list_desc,
-                sota_exp_desc=sota_exp_desc,
-            )
-            for problem_name in fb_problems:
-                fb_problems[problem_name]["label"] = "FEEDBACK_PROBLEM"
-                all_problems[problem_name] = fb_problems[problem_name]
 
-        if len(trace.hist) < 9:
-            scen_problems = self.identify_scenario_problem(
-                scenario_desc=scenario_desc,
-                sota_exp_desc=sota_exp_desc,
-            )
-            for problem_name in scen_problems:
-                scen_problems[problem_name]["label"] = "SCENARIO_PROBLEM"
-                all_problems[problem_name] = scen_problems[problem_name]
+        all_problems = self.identify_problems_v3(
+            trace=trace,
+            scenario_desc=scenario_desc,
+            sota_exp_desc=sota_exp_desc,
+            exp_feedback_list_desc=exp_feedback_list_desc,
+        )
+
+        # if len(trace.hist) > 3:
+        #     fb_problems = self.identify_feedback_problem(
+        #         scenario_desc=scenario_desc,
+        #         exp_feedback_list_desc=exp_feedback_list_desc,
+        #         sota_exp_desc=sota_exp_desc,
+        #     )
+        #     for problem_name in fb_problems:
+        #         fb_problems[problem_name]["label"] = "FEEDBACK_PROBLEM"
+        #         all_problems[problem_name] = fb_problems[problem_name]
+
+        # if len(trace.hist) < 9:
+        #     scen_problems = self.identify_scenario_problem(
+        #         scenario_desc=scenario_desc,
+        #         sota_exp_desc=sota_exp_desc,
+        #     )
+        #     for problem_name in scen_problems:
+        #         scen_problems[problem_name]["label"] = "SCENARIO_PROBLEM"
+        #         all_problems[problem_name] = scen_problems[problem_name]
 
         # Step 1.5: Sample ideas from idea pool
         if DS_RD_SETTING.enable_knowledge_base:
@@ -1128,7 +1185,6 @@ def gen(self, trace: DSTrace) -> DSExperiment:
         pickled_problem_name, new_hypothesis = self.hypothesis_rank(
             hypothesis_dict=hypothesis_dict,
             problem_dict=all_problems,
-            trace=trace,
         )
         # Step 3.5: Update knowledge base with the picked problem
         if DS_RD_SETTING.enable_knowledge_base:

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

@@ -47,7 +47,7 @@ def get_sota_exp_to_submit(self, trace: Trace) -> DSExperiment | None:
         sota_exp_fb_list = trace.experiment_and_feedback_list_after_init(
             return_type="sota", search_type="all", max_retrieve_num=DS_RD_SETTING.max_sota_retrieved_num
         )
-
+        logger.info(f"Auto SOTA selector: Found {len(sota_exp_fb_list)} SOTA experiments")
         if len(sota_exp_fb_list) == 0:
             logger.info("Auto SOTA selector: No SOTA in trace yet")
             return None
@@ -69,6 +69,8 @@ def get_sota_exp_to_submit(self, trace: Trace) -> DSExperiment | None:
             leaves: list[int] = trace.get_leaves()
 
             if len(leaves) >= 2:
+
+                logger.info(f"Auto SOTA selector: Multiple traces found, collecting SOTA experiments from each trace")
                 # multiple trace case, collect the latest SOTA experiments from each trace
                 new_sota_exp_fb_list: list[tuple[DSExperiment, ExperimentFeedback]] = []
                 # calculate the number of SOTA experiments to retrieve from each trace
@@ -81,11 +83,26 @@ def get_sota_exp_to_submit(self, trace: Trace) -> DSExperiment | None:
                         selection=(leaf,),
                         max_retrieve_num=max_sota_retrieved_num_per_trace,
                     )
+                    logger.info(
+                        f"Auto SOTA selector: Collected {len(sota_exp_fb_list_per_trace)} SOTA experiments from trace with leaf #. {leaf}"
+                    )
 
                     new_sota_exp_fb_list.extend(sota_exp_fb_list_per_trace)
 
                 sota_exp_fb_list = new_sota_exp_fb_list
 
+                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:
+                    logger.info("Auto SOTA selector: Only one SOTA in trace, using it")
+                    return sota_exp_fb_list[0][0]
+                else:
+                    logger.info(
+                        f"Auto SOTA selector: {len(sota_exp_fb_list)} SOTA experiments found in all traces, calling LLM to select the best one"
+                    )
+
             for i, (exp, ef) in enumerate(sota_exp_fb_list):
                 if exp:
                     current_final_score = pd.DataFrame(exp.result).loc["ensemble"].iloc[0]
@@ -115,7 +132,7 @@ def get_sota_exp_to_submit(self, trace: Trace) -> DSExperiment | None:
 
             sota_submit_idx = response_dict.get("selected_SOTA_idx", None)
 
-            if sota_submit_idx is not None:
+            if sota_submit_idx and int(sota_submit_idx) - 1 < len(sota_exp_fb_list):
                 sota_submit = sota_exp_fb_list[int(sota_submit_idx) - 1]
                 sota_idx_in_trace = trace.hist.index(sota_submit)
                 logger.info(
@@ -124,8 +141,12 @@ def get_sota_exp_to_submit(self, trace: Trace) -> DSExperiment | None:
                 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]
+                if len(sota_exp_fb_list) > 0:
+                    logger.info("Auto SOTA selector: No SOTA experiment to submit, using the latest SOTA experiment")
+                    return sota_exp_fb_list[-1][0]
+                else:
+                    logger.info("Auto SOTA selector: No SOTA experiment in trace yet")
+                    return None
 
 
 class BestValidSelector(SOTAexpSelector):