feat: enable meta planner (#1103)

* enable meta planner

* fix a small bug

* ADD PLAN TO GEN

* remove ensemble in planner

* fix CI

* fix CI

* fix planner threshold

rdagent/app/data_science/conf.py

@@ -18,7 +18,8 @@ class DataScienceBasePropSetting(KaggleBasePropSetting):
     - For custom data science scenarios, use: "rdagent.scenarios.data_science.scen.DataScienceScen"
     """
 
-    hypothesis_gen: str = "rdagent.scenarios.data_science.proposal.exp_gen.proposal.DSProposalV2ExpGen"
+    planner: str = "rdagent.scenarios.data_science.proposal.exp_gen.planner.DSExpPlannerHandCraft"
+    hypothesis_gen: str = "rdagent.scenarios.data_science.proposal.exp_gen.router.ParallelMultiTraceExpGen"
     """Hypothesis generation class"""
 
     summarizer: str = "rdagent.scenarios.data_science.dev.feedback.DSExperiment2Feedback"
@@ -99,15 +100,12 @@ class DataScienceBasePropSetting(KaggleBasePropSetting):
     # inject diverse when start a new sub-trace
     enable_inject_diverse: bool = False
 
-    # 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
+    merge_hours: int = 0
     """The time for merge"""
 
     #### multi-trace: max SOTA-retrieved number, used in AutoSOTAexpSelector
@@ -115,5 +113,11 @@ class DataScienceBasePropSetting(KaggleBasePropSetting):
     max_sota_retrieved_num: int = 10
     """The maximum number of SOTA experiments to retrieve in a LLM call"""
 
+    #### enable draft before first sota experiment
+    enable_draft_before_first_sota: bool = False
+    enable_planner: bool = False
+
+    model_architecture_suggestion_time_percent: float = 0.75
+
 
 DS_RD_SETTING = DataScienceBasePropSetting()

rdagent/components/proposal/__init__.py

@@ -3,6 +3,7 @@ from typing import Tuple
 
 from rdagent.core.experiment import Experiment
 from rdagent.core.proposal import (
+    ExperimentPlan,
     Hypothesis,
     Hypothesis2Experiment,
     HypothesisGen,
@@ -25,7 +26,11 @@ class LLMHypothesisGen(HypothesisGen):
     @abstractmethod
     def convert_response(self, response: str) -> Hypothesis: ...
 
-    def gen(self, trace: Trace) -> Hypothesis:
+    def gen(
+        self,
+        trace: Trace,
+        plan: ExperimentPlan | None = None,
+    ) -> Hypothesis:
         context_dict, json_flag = self.prepare_context(trace)
 
         system_prompt = T(".prompts:hypothesis_gen.system_prompt").r(

rdagent/core/experiment.py

@@ -292,6 +292,12 @@ ASpecificWSForExperiment = TypeVar("ASpecificWSForExperiment", bound=Workspace)
 ASpecificWSForSubTasks = TypeVar("ASpecificWSForSubTasks", bound=Workspace)
 
 
+class ExperimentPlan(dict[str, Any]):
+    """
+    A plan for the experiment, which is a dictionary that contains the plan to each stage.
+    """
+
+
 class Experiment(
     ABC,
     Generic[ASpecificTask, ASpecificWSForExperiment, ASpecificWSForSubTasks],
@@ -337,6 +343,9 @@ class Experiment(
 
         # For parallel multi-trace support
         self.local_selection: tuple[int, ...] | None = None
+        self.plan: ExperimentPlan | None = (
+            None  # To store the planning information for this experiment, should be generated inside exp_gen.gen
+        )
 
     @property
     def result(self) -> object:
@@ -348,6 +357,7 @@ class Experiment(
 
 
 ASpecificExp = TypeVar("ASpecificExp", bound=Experiment)
+ASpecificPlan = TypeVar("ASpecificPlan", bound=ExperimentPlan)
 
 TaskOrExperiment = TypeVar("TaskOrExperiment", Task, Experiment)
 

rdagent/core/proposal.py

@@ -8,7 +8,12 @@ from typing import TYPE_CHECKING, Generic, TypeVar
 
 from rdagent.core.conf import RD_AGENT_SETTINGS
 from rdagent.core.evaluation import Feedback
-from rdagent.core.experiment import ASpecificExp, Experiment
+from rdagent.core.experiment import (
+    ASpecificExp,
+    ASpecificPlan,
+    Experiment,
+    ExperimentPlan,
+)
 from rdagent.core.knowledge_base import KnowledgeBase
 from rdagent.core.scenario import Scenario
 
@@ -268,15 +273,34 @@ class SOTAexpSelector:
         """
 
 
+class ExpPlanner(ABC, Generic[ASpecificPlan]):
+    """
+    An abstract class for planning the experiment.
+    The planner should generate a plan for the experiment based on the trace.
+    """
+
+    def __init__(self, scen: Scenario) -> None:
+        self.scen = scen
+
+    @abstractmethod
+    def plan(self, trace: Trace) -> ASpecificPlan:
+        """
+        Generate a plan for the experiment based on the trace.
+        The plan should be a dictionary that contains the plan to each stage.
+        """
+
+
 class ExpGen(ABC):
 
     def __init__(self, scen: Scenario) -> None:
         self.scen = scen
 
     @abstractmethod
-    def gen(self, trace: Trace) -> Experiment:
+    def gen(self, trace: Trace, plan: ExperimentPlan | None = None) -> Experiment:
         """
         Generate the experiment based on the trace.
+        Planning is part of gen, but since we may support multi-stage planning,
+        we need to pass plan as optional argument.
 
         `ExpGen().gen()` play a role like
 
@@ -306,7 +330,11 @@ class HypothesisGen(ABC):
         self.scen = scen
 
     @abstractmethod
-    def gen(self, trace: Trace) -> Hypothesis:
+    def gen(
+        self,
+        trace: Trace,
+        plan: ExperimentPlan | None = None,
+    ) -> Hypothesis:
         # def gen(self, scenario_desc: str, ) -> Hypothesis:
         """
         Motivation of the variable `scenario_desc`:

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

@@ -15,6 +15,7 @@ from rdagent.log import rdagent_logger as logger
 from rdagent.oai.llm_utils import APIBackend
 from rdagent.scenarios.data_science.experiment.experiment import COMPONENT, DSExperiment
 from rdagent.scenarios.data_science.proposal.exp_gen.base import DSHypothesis, DSTrace
+from rdagent.scenarios.data_science.proposal.exp_gen.planner import DSExperimentPlan
 from rdagent.scenarios.data_science.proposal.exp_gen.utils import (
     CodingSketch,
     get_component,
@@ -61,6 +62,7 @@ class DSDraftExpGen(ExpGen):
         self,
         component: COMPONENT,
         trace: DSTrace,
+        plan: DSExperimentPlan | None = None,
     ) -> DSExperiment:
         """Handle any component using a unified approach.
 
@@ -234,7 +236,11 @@ class DSDraftV2ExpGen(ExpGen):
             exp.pending_tasks_list.append([workflow_task])
         return exp
 
-    def gen(self, trace: DSTrace) -> DSExperiment:
+    def gen(
+        self,
+        trace: DSTrace,
+        plan: DSExperimentPlan | None = None,
+    ) -> DSExperiment:
         # Step 0: Prepare
         pipeline = DS_RD_SETTING.coder_on_whole_pipeline
         if pipeline:

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

@@ -13,13 +13,18 @@ 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.planner import DSExperimentPlan
 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):
-    def gen(self, trace: DSTrace) -> DSExperiment:
+    def gen(
+        self,
+        trace: DSTrace,
+        plan: DSExperimentPlan | None = None,
+    ) -> DSExperiment:
         # Ignore the selection argument and use all leaves instead.
         leaves: list[int] = trace.get_leaves()
         trace.set_current_selection((leaves[0],))  # override the current selection.
@@ -136,7 +141,11 @@ class ExpGen2Hypothesis(DSProposalV2ExpGen):
                 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:
+    def gen(
+        self,
+        trace: DSTrace,
+        plan: DSExperimentPlan | None = None,
+    ) -> DSExperiment:
         # Ignore the selection argument and use all leaves instead.
         sota_exp_fb = trace.sota_experiment_fb(selection=trace.current_selection)
 
@@ -231,7 +240,11 @@ class ExpGen2TraceAndMerge(ExpGen):
         self.merge_exp_gen = MergeExpGen(self.scen)
         self.exp_gen = DataScienceRDLoop.default_exp_gen(self.scen)
 
-    def gen(self, trace: DSTrace) -> DSExperiment:
+    def gen(
+        self,
+        trace: DSTrace,
+        plan: DSExperimentPlan | None = None,
+    ) -> DSExperiment:
         timer: RDAgentTimer = RD_Agent_TIMER_wrapper.timer
         logger.info(f"Remain time: {timer.remain_time()}")
 
@@ -257,7 +270,11 @@ class ExpGen2TraceAndMerge(ExpGen):
 
 
 class MergeExpGen_MultiTrace(ExpGen):
-    def gen(self, trace: DSTrace) -> DSExperiment:
+    def gen(
+        self,
+        trace: DSTrace,
+        plan: DSExperimentPlan | None = None,
+    ) -> DSExperiment:
         # Ignore the selection argument and use all leaves instead.
         leaves: list[int] = trace.get_leaves()
 
@@ -347,18 +364,13 @@ class ExpGen2TraceAndMergeV2(ExpGen):
         # )
         raise NotImplementedError("You should not switch version with proposal_version")
 
-    def gen(self, trace: DSTrace, selection: tuple[int, ...] = (-1,)) -> DSExperiment:
+    def gen(
+        self, trace: DSTrace, plan: DSExperimentPlan | None = None, selection: tuple[int, ...] = (-1,)
+    ) -> DSExperiment:
         timer: RDAgentTimer = RD_Agent_TIMER_wrapper.timer
         logger.info(f"Remain time: {timer.remain_time()}")
 
         if timer.remain_time() >= 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:
@@ -402,21 +414,15 @@ class ExpGen2TraceAndMergeV3(ExpGen):
         self.merge_exp_gen = ExpGen2Hypothesis(self.scen)
         self.exp_gen = DataScienceRDLoop.default_exp_gen(self.scen)
 
-    def gen(self, trace: DSTrace) -> DSExperiment:
+    def gen(
+        self,
+        trace: DSTrace,
+        plan: DSExperimentPlan | None = None,
+    ) -> DSExperiment:
         timer: RDAgentTimer = RD_Agent_TIMER_wrapper.timer
         logger.info(f"Remain time: {timer.remain_time()}")
 
         if timer.remain_time() >= 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

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

@@ -6,12 +6,17 @@ from rdagent.components.coder.data_science.pipeline.exp import PipelineTask
 from rdagent.core.proposal import ExpGen
 from rdagent.scenarios.data_science.experiment.experiment import DSExperiment
 from rdagent.scenarios.data_science.proposal.exp_gen.base import DSHypothesis, DSTrace
+from rdagent.scenarios.data_science.proposal.exp_gen.router import DSExperimentPlan
 from rdagent.utils.agent.tpl import T
 from rdagent.utils.agent.workflow import build_cls_from_json_with_retry
 
 
 class NaiveExpGen(ExpGen):
-    def gen(self, trace: DSTrace) -> DSExperiment:
+    def gen(
+        self,
+        trace: DSTrace,
+        plan: DSExperimentPlan | None = None,
+    ) -> DSExperiment:
         sota_exp = trace.sota_experiment()
         scenario_desc = trace.scen.get_scenario_all_desc()
         sota_exp_desc = T("scenarios.data_science.share:describe.exp").r(

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

@@ -6,7 +6,7 @@ from typing import TYPE_CHECKING
 
 from rdagent.app.data_science.conf import DS_RD_SETTING
 from rdagent.core.conf import RD_AGENT_SETTINGS
-from rdagent.core.proposal import ExpGen
+from rdagent.core.proposal import ExperimentPlan, ExpGen
 from rdagent.log import rdagent_logger as logger
 from rdagent.log.timer import RD_Agent_TIMER_wrapper, RDAgentTimer
 from rdagent.scenarios.data_science.loop import DataScienceRDLoop
@@ -38,7 +38,11 @@ class ParallelMultiTraceExpGen(ExpGen):
         self.merge_exp_gen = ExpGen2Hypothesis(self.scen)
         self.trace_scheduler: TraceScheduler = RoundRobinScheduler(DS_RD_SETTING.max_trace_num)
 
-    def gen(self, trace: "DSTrace") -> "Experiment":
+    def gen(
+        self,
+        trace: "DSTrace",
+        plan: "ExperimentPlan" | None = None,
+    ) -> "Experiment":
         raise NotImplementedError(
             "ParallelMultiTraceExpGen is designed for async usage, please call async_gen instead."
         )
@@ -57,16 +61,6 @@ class ParallelMultiTraceExpGen(ExpGen):
 
             if timer.remain_time() >= 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
-
                 if loop.get_unfinished_loop_cnt(loop.loop_idx) < RD_AGENT_SETTINGS.get_max_parallel():
                     local_selection = await self.trace_scheduler.next(trace)
 

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

@@ -0,0 +1,45 @@
+from datetime import timedelta
+
+from rdagent.app.data_science.conf import DS_RD_SETTING
+from rdagent.components.coder.CoSTEER import RD_Agent_TIMER_wrapper
+from rdagent.core.proposal import ExperimentPlan, ExpPlanner
+from rdagent.scenarios.data_science.proposal.exp_gen.base import DSTrace
+
+
+class DSExperimentPlan(ExperimentPlan):
+    """
+    A specific plan for data science experiments.
+    This plan can include various stages such as proposal, draft, and merge.
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.setdefault("exp_gen", {}).setdefault("draft", False)
+        self.setdefault("exp_gen", {}).setdefault("suggest_model_architecture", False)
+        self.setdefault("exp_gen", {}).setdefault("suggest_model_ensemble", False)
+
+
+class DSExpPlannerHandCraft(ExpPlanner[DSExperimentPlan]):
+    """
+    A specific planner for data science experiments.
+    """
+
+    def plan(self, trace: DSTrace) -> DSExperimentPlan:
+        """
+        Generate a plan for the experiment based on the trace.
+        The plan should be a dictionary that contains the plan to each stage.
+        trace is well selected into sub trace mode
+        """
+        plan = DSExperimentPlan()
+        timer = RD_Agent_TIMER_wrapper.timer
+        remain_percent = timer.remain_time() / timer.all_duration if timer.started else 1.0
+
+        if not trace.sota_experiment():
+            plan["exp_gen"]["draft"] = True
+        elif trace.sota_experiment() and remain_percent > DS_RD_SETTING.model_architecture_suggestion_time_percent:
+            plan["exp_gen"]["suggest_model_architecture"] = True
+        # elif DS_RD_SETTING.merge_hours > 0:
+        #     merge_percent = timedelta(hours=DS_RD_SETTING.merge_hours) / timer.all_duration
+        #     if merge_percent < remain_percent < merge_percent + 0.1:
+        #         plan["exp_gen"]["suggest_model_ensemble"] = True
+        return plan

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

@@ -10,10 +10,11 @@ scenario_problem:
     Your task is to analyze the provided information (primarily the scenario and current SOTA, if available) and identify a concise list of **Key Challenges** or **Core Problems** relevant to achieving success in this competition and improving the target metric. Aim for **FEWER BUT BETTER** challenges (e.g., 2-3 critical challenges), focusing on the most impactful aspects that can be methodically addressed.
 
     ### Core Analysis Dimensions for Identifying Challenges
-    1. **SOTA Alignment Analysis**: (If SOTA is provided) Systematically compare the current SOTA implementation against dataset properties and domain knowledge to identify discrepancies or areas representing core challenges to overcome for enhancement.
-    2. **Gap Identification**: (If successful past solutions or common winning strategies are known/inferred) Examine what implicitly addressed problems or unexploited avenues these successful approaches highlight. These gaps can represent current challenges.
-    3. **Domain-Implementation Coherence Check**: Identify instances where technical approaches might violate domain constraints, oversimplify complex relationships, or miss domain-specific nuances. These incoherencies are challenges.
-    4. **Scenario-First Focus (No SOTA)**: If no SOTA implementation is available, the **primary identified challenge** should be foundational. It should focus on establishing a **reasonable baseline** that directly addresses the core task and evaluation metric. Avoid overly complex initial challenges.
+    - **Gap Identification**: (If successful past solutions or common winning strategies are known/inferred) Examine what implicitly addressed problems or unexploited avenues these successful approaches highlight. These gaps can represent current challenges.
+    - **Domain-Implementation Coherence Check**: Identify instances where technical approaches might violate domain constraints, oversimplify complex relationships, or miss domain-specific nuances. These incoherencies are challenges.
+    {% if plan.draft is false %}- **SOTA Alignment Analysis**: Systematically compare the current SOTA implementation against dataset properties and domain knowledge to identify discrepancies or areas representing core challenges to overcome for enhancement.
+    {% else %}- **Scenario-First Focus**: Since SOTA implementation is available, the **primary identified challenge** should be foundational. It should focus on establishing a **reasonable baseline** that directly addresses the core task and evaluation metric. Avoid overly complex initial challenges.
+    {% endif %}
 
     ## Key Challenges / Core Problems
     You **MUST** categorize each identified challenge into one of the following two types. This categorization should be based on the primary driver or nature of the challenge:
@@ -24,7 +25,7 @@ scenario_problem:
     1. The challenge should be specific and fine-grained. Avoid general or vague statements.
     2. The challenge should be technical or methodological. Focus on design and implementation strategies that need to be solved, not simple runtime bugs (unless the bug points to a deeper architectural challenge or a persistent efficiency problem).
     3. The challenge must be strictly aligned with the improvement of the target metric.
-    4. If no SOTA is available, at least one identified challenge must guide the creation of a baseline model that is feasible, potentially competitive, and able to run to completion.
+    {% if plan.draft is true %}4. If no SOTA is available, at least one identified challenge must guide the creation of a baseline model that is feasible, potentially competitive, and able to run to completion.{% endif %}
 
 
     {% if problem_output_format is not none %}
@@ -54,7 +55,6 @@ feedback_problem:
 
     ## Key Learnings and Unresolved Challenges
 
-    
     {% 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.
@@ -137,10 +137,12 @@ hypothesis_gen:
       - Analyze the Identified Challenge to understand its root cause and potential impact on the competition's target metric or successful execution.
       - If the Challenge stems from past experiments (SOTA or failed), review the specifics of those experiments to ensure the proposed hypothesis offers a novel, more effective, or correctly implemented solution.
       - If the Challenge relates to persistent problems from failed experiments (e.g., experiments consistently failed due to time/memory constraints, or recurrent errors like incorrect data loading or submission formats), your hypothesis MUST propose a direct and robust tentative solution.
+    {% if plan.draft is true %}
     2. **Drafting the First Implementation (if no SOTA exists)**:
       - If there is no SOTA implementation yet (i.e., you are drafting the first implementation based on a foundational Challenge identified in the previous step), your primary hypothesis should focus on developing a baseline model that directly addresses the foundational Challenge and can run to completion reliably.
       - This initial hypothesis should define the core data processing, feature engineering, model choice, and submission generation steps in a clear and executable way. Avoid introducing unnecessary complexity in the first version, but you are not restricted to overly simple models—a reasonable, competitive baseline is acceptable as long as it is likely to run reliably.
-    3. **Actionable Changes**:
+    {% endif %}
+    {% if plan.draft is true %}3{% else %}2{% endif %}. **Actionable Changes**:
       - If a Challenge involves underperforming models (e.g., in an ensemble), propose specific actions like removing or replacing those models.
       - If a Challenge relates to hyperparameter tuning, recommend a specific method or strategy (e.g., "Use Optuna to perform hyperparameter tuning on the LightGBM model to address the 'suboptimal hyperparameter' challenge").
       - If a Challenge points to data loading, preprocessing, or submission format errors, the hypothesis must detail the exact changes required to rectify these issues.
@@ -199,6 +201,17 @@ hypothesis_gen:
     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 %}
+
+    {% if plan.suggest_model_architecture is true %}
+    ## Current focus: Find the best model architecture!
+    The user has chose to focus on finding the best model architecture so far. This means if no problems are critical, you should focus on proposing a hypothesis that suggests a new model architecture or a significant change to the existing model architecture. This is the primary focus of the current iteration.
+    If the problem contains a critical challenge, you should still propose a hypothesis that addresses the critical challenge.
+    {% elif plan.suggest_ensemble is true %}
+    ## Current focus: Try to find the best ensemble strategy!
+    The user has chose to focus on finding the best ensemble strategy so far. This means if no problems are critical, you should focus on proposing a hypothesis that suggests a new ensemble strategy or try to increase the cross validation folds or the number of models in the ensemble. This is the primary focus of the current iteration.
+    Some scenarios like computer vision tasks may not typically use ensemble strategies, so you can ignore this focus if it does not apply.
+    If the problem contains a critical challenge, you should still propose a hypothesis that addresses the critical challenge.
+    {% endif %}
     
     {% if hypothesis_output_format is not none %}
     ## Final Output Format in JSON Schema:

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

@@ -23,6 +23,7 @@ from rdagent.scenarios.data_science.proposal.exp_gen.draft.draft import (
     DSDraftExpGen,  # TODO: DSDraftExpGen should be moved to router in the further
 )
 from rdagent.scenarios.data_science.proposal.exp_gen.idea_pool import DSIdea
+from rdagent.scenarios.data_science.proposal.exp_gen.planner import DSExperimentPlan
 from rdagent.scenarios.data_science.proposal.exp_gen.utils import get_packages
 from rdagent.utils.agent.tpl import T
 from rdagent.utils.repo.diff import generate_diff_from_dict
@@ -294,7 +295,11 @@ def draft_exp_in_decomposition(scen: Scenario, trace: DSTrace) -> None | DSDraft
 
 
 class DSProposalV1ExpGen(ExpGen):
-    def gen(self, trace: DSTrace) -> DSExperiment:
+    def gen(
+        self,
+        trace: DSTrace,
+        plan: DSExperimentPlan | None = None,
+    ) -> DSExperiment:
         # Drafting Stage
         if draft_exp := draft_exp_in_decomposition(self.scen, trace):
             return draft_exp
@@ -467,11 +472,17 @@ class DSProposalV2ExpGen(ExpGen):
         super().__init__(*args, **kwargs)
         self.supports_response_schema = APIBackend().supports_response_schema()
 
-    def identify_scenario_problem(self, scenario_desc: str, sota_exp_desc: str) -> Dict:
+    def identify_scenario_problem(
+        self,
+        scenario_desc: str,
+        sota_exp_desc: str,
+        exp_gen_plan: Dict,
+    ) -> Dict:
         sys_prompt = T(".prompts_v2:scenario_problem.system").r(
             problem_output_format=(
                 T(".prompts_v2:output_format.problem").r() if not self.supports_response_schema else None
             ),
+            plan=exp_gen_plan,
         )
         user_prompt = T(".prompts_v2:scenario_problem.user").r(
             scenario_desc=scenario_desc,
@@ -524,7 +535,13 @@ class DSProposalV2ExpGen(ExpGen):
         return problems
 
     def identify_problem(
-        self, current_sub_trace, scenario_desc, sota_exp_desc, exp_feedback_list_desc, inject_diverse
+        self,
+        current_sub_trace,
+        scenario_desc,
+        sota_exp_desc,
+        exp_feedback_list_desc,
+        inject_diverse,
+        exp_gen_plan,
     ) -> Dict:
         sota_exp_num = sum(1 for _, fb in current_sub_trace if fb.decision)
         failed_exp_num = len(current_sub_trace) - sota_exp_num
@@ -536,6 +553,7 @@ class DSProposalV2ExpGen(ExpGen):
             scen_problems = self.identify_scenario_problem(
                 scenario_desc=scenario_desc,
                 sota_exp_desc=sota_exp_desc,
+                exp_gen_plan=exp_gen_plan,
             )
             for problem_name in scen_problems:
                 scen_problems[problem_name]["label"] = "SCENARIO_PROBLEM"
@@ -564,6 +582,7 @@ class DSProposalV2ExpGen(ExpGen):
         pipeline: bool,
         enable_idea_pool: bool,
         inject_diverse: bool = False,
+        exp_gen_plan: Optional[Dict] = None,
     ) -> Dict:
         problem_formatted_str = ""
         for i, (problem_name, problem_dict) in enumerate(problems.items()):
@@ -583,6 +602,7 @@ class DSProposalV2ExpGen(ExpGen):
             pipeline=pipeline,
             enable_idea_pool=enable_idea_pool,
             inject_diverse=inject_diverse,
+            plan=exp_gen_plan,
         )
         user_prompt = T(".prompts_v2:hypothesis_gen.user").r(
             scenario_desc=scenario_desc,
@@ -822,6 +842,7 @@ class DSProposalV2ExpGen(ExpGen):
     def gen(
         self,
         trace: DSTrace,
+        plan: DSExperimentPlan | None = None,
     ) -> DSExperiment:
         pipeline = DS_RD_SETTING.coder_on_whole_pipeline
         if not pipeline and (draft_exp := draft_exp_in_decomposition(self.scen, trace)):
@@ -881,6 +902,7 @@ class DSProposalV2ExpGen(ExpGen):
             sota_exp_desc=sota_exp_desc,
             exp_feedback_list_desc=exp_feedback_list_desc,
             inject_diverse=inject_diverse,
+            exp_gen_plan=plan.get("exp_gen") if plan else None,
         )
 
         # Step 1.5: Sample ideas from idea pool
@@ -903,6 +925,7 @@ class DSProposalV2ExpGen(ExpGen):
             pipeline=pipeline,
             enable_idea_pool=DS_RD_SETTING.enable_knowledge_base,
             inject_diverse=inject_diverse,
+            exp_gen_plan=plan.get("exp_gen") if plan else None,
         )
         if not pipeline:
             sota_exp_model_file_count = len(

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

@@ -1,25 +1,116 @@
+from __future__ import annotations
+
+import asyncio
+from datetime import timedelta
+from typing import TYPE_CHECKING
+
 from rdagent.app.data_science.conf import DS_RD_SETTING
-from rdagent.core.proposal import ExpGen
+from rdagent.core.conf import RD_AGENT_SETTINGS
+from rdagent.core.proposal import ExpGen, ExpPlanner
+from rdagent.core.utils import import_class
+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.loop import DataScienceRDLoop
 from rdagent.scenarios.data_science.proposal.exp_gen.base import DSTrace
 from rdagent.scenarios.data_science.proposal.exp_gen.draft.draft import DSDraftV2ExpGen
+from rdagent.scenarios.data_science.proposal.exp_gen.merge import ExpGen2Hypothesis
+from rdagent.scenarios.data_science.proposal.exp_gen.planner import (
+    DSExperimentPlan,
+    ExperimentPlan,
+)
 from rdagent.scenarios.data_science.proposal.exp_gen.proposal import DSProposalV2ExpGen
+from rdagent.scenarios.data_science.proposal.exp_gen.trace_scheduler import (
+    RoundRobinScheduler,
+    TraceScheduler,
+)
+
+if TYPE_CHECKING:
+    from rdagent.scenarios.data_science.experiment.experiment import DSExperiment
+    from rdagent.scenarios.data_science.proposal.exp_gen.base import DSTrace, Experiment
+    from rdagent.utils.workflow.loop import LoopBase
 
 
-class DraftRouterExpGen(ExpGen):
+class ParallelMultiTraceExpGen(ExpGen):
     """
-    A intelligent router for drafting and proposing.
+    An experiment generation strategy that enables parallel multi-trace exploration.
 
+    This generator is designed to work with the "Attribute Injection" model.
+    It uses a TraceScheduler to determine which parent node to expand, and
+    injects this parent context into the experiment object itself.
     """
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
+        # The underlying generator for creating a single experiment
+        self.exp_gen = DataScienceRDLoop.default_exp_gen(self.scen)
         self.draft_exp_gen = DSDraftV2ExpGen(self.scen)
-        self.base_exp_gen = DSProposalV2ExpGen(self.scen)
+        self.merge_exp_gen = ExpGen2Hypothesis(self.scen)
+        self.trace_scheduler: TraceScheduler = RoundRobinScheduler(DS_RD_SETTING.max_trace_num)
+        self.planner = import_class(DS_RD_SETTING.planner)(self.scen)
 
-    def gen(self, trace: DSTrace) -> DSExperiment:
-        pipeline = DS_RD_SETTING.coder_on_whole_pipeline
-        sota_exp = trace.sota_experiment()
-        if sota_exp is None and pipeline:
-            return self.draft_exp_gen.gen(trace)
-        return self.base_exp_gen.gen(trace)
+    def gen(
+        self,
+        trace: "DSTrace",
+        plan: "ExperimentPlan" | None = None,
+    ) -> "Experiment":
+        raise NotImplementedError(
+            "ParallelMultiTraceExpGen is designed for async usage, please call async_gen instead."
+        )
+
+    async def async_gen(self, trace: DSTrace, loop: LoopBase) -> DSExperiment:
+        """
+        Waits for a free execution slot, selects a parent trace using the
+        scheduler, generates a new experiment, and injects the parent context
+        into it before returning.
+        """
+        timer: RDAgentTimer = RD_Agent_TIMER_wrapper.timer
+        logger.info(f"Remain time: {timer.remain_time()}")
+        local_selection: tuple[int, ...] = None
+
+        while True:
+            if loop.get_unfinished_loop_cnt(loop.loop_idx) < RD_AGENT_SETTINGS.get_max_parallel():
+                # set trace current selection
+                if not timer.started or timer.remain_time() >= timedelta(hours=DS_RD_SETTING.merge_hours):
+                    local_selection = await self.trace_scheduler.next(trace)
+
+                    # set the local selection as the global current selection for the trace
+                    trace.set_current_selection(local_selection)
+                else:
+                    leaves: list[int] = trace.get_leaves()
+                    if len(leaves) < 2:
+                        local_selection = (-1,)
+                        trace.set_current_selection(selection=local_selection)
+                    else:
+                        local_selection = (leaves[0],)
+                        if trace.sota_exp_to_submit is not None:
+                            for i in range(1, len(leaves)):
+                                if trace.is_parent(trace.exp2idx(trace.sota_exp_to_submit), leaves[i]):
+                                    local_selection = (leaves[i],)
+                                    break
+                        trace.set_current_selection(local_selection)
+
+                ds_plan = self.planner.plan(trace) if DS_RD_SETTING.enable_planner else DSExperimentPlan()
+                if (
+                    (not timer.started or timer.remain_time() >= timedelta(hours=DS_RD_SETTING.merge_hours))
+                    and trace.sota_experiment(selection=local_selection) is None
+                    and DS_RD_SETTING.enable_draft_before_first_sota
+                ):
+                    exp = self.draft_exp_gen.gen(trace, plan=ds_plan)
+                elif (
+                    timer.started
+                    and timer.remain_time() < timedelta(hours=DS_RD_SETTING.merge_hours)
+                    and len(leaves) >= 2
+                ):
+                    DS_RD_SETTING.coding_fail_reanalyze_threshold = 100000
+                    DS_RD_SETTING.consecutive_errors = 100000
+                    exp = self.merge_exp_gen.gen(trace, plan=ds_plan)
+                else:
+                    # If there is a sota experiment in the sub-trace and not in merge time, we use default exp_gen
+                    exp = self.exp_gen.gen(trace, plan=ds_plan)
+
+                exp.set_local_selection(local_selection)
+                exp.plan = ds_plan
+                return exp
+
+            await asyncio.sleep(1)

rdagent/scenarios/data_science/share.yaml

@@ -52,7 +52,10 @@ describe: # some template to describe some object
     {% if not pipeline %}Chosen Component: {{ exp_and_feedback[0].hypothesis.component }}{% endif %}
     Proposed Hypothesis: {{ exp_and_feedback[0].hypothesis.hypothesis }}
     {% if exp_and_feedback[1].code_change_summary  %}Code Change Summary: {{ exp_and_feedback[1].code_change_summary }}{% endif %}
-    Surpass Previous SOTA: {{ exp_and_feedback[1].decision }}
+    Surpass Previous SOTA: {{ exp_and_feedback[1].decision }}    
+    {% if exp_and_feedback[0].running_info.running_time is not none %}
+    Experiment Running Time: {{ exp_and_feedback[0].running_info.running_time }} seconds
+    {% endif %}
     {% if exp_and_feedback[0].result is none %}
     Experiment Score: Running buggy
     Experiment Error: {{ exp_and_feedback[1].reason }}