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

Author: peteryang1

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,21 +100,24 @@ 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
     # constrains the number of SOTA experiments to retrieve, otherwise too many SOTA experiments to retrieve will cause the exceed of the context window of LLM
     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 rdagent.core.experiment import Experiment
 from rdagent.core.proposal import (
+    ExperimentPlan,
     Hypothesis,
     Hypothesis2Experiment,
     HypothesisGen,
@@ -25,7 +26,11 @@ def prepare_context(self, trace: Trace) -> Tuple[dict, bool]: ...
     @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 @@ def __str__(self) -> str:
 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 @@ def __init__(
 
         # 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 @@ def result(self, value: object) -> None:
 
 
 ASpecificExp = TypeVar("ASpecificExp", bound=Experiment)
+ASpecificPlan = TypeVar("ASpecificPlan", bound=ExperimentPlan)
 
 TaskOrExperiment = TypeVar("TaskOrExperiment", Task, Experiment)
 

rdagent/core/proposal.py

@@ -8,7 +8,12 @@
 
 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 @@ def get_sota_exp_to_submit(self, trace: Trace) -> Experiment | None:
         """
 
 
+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 @@ def __init__(self, scen: Scenario) -> None:
         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.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 @@ def gen(
         self,
         component: COMPONENT,
         trace: DSTrace,
+        plan: DSExperimentPlan | None = None,
     ) -> DSExperiment:
         """Handle any component using a unified approach.
 
@@ -234,7 +236,11 @@ def task_gen(
             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.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 @@ def get_exp_index(self, trace: DSTrace) -> int:
                 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 @@ def __init__(self, *args, **kwargs):
         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 @@ def gen(self, trace: DSTrace) -> DSExperiment:
 
 
 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 @@ def reset_exp_gen_version(self, version: str = "v2"):
         # )
         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 @@ def __init__(self, *args, **kwargs):
         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.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 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 @@ def __init__(self, *args, **kwargs):
         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 @@ async def async_gen(self, trace: DSTrace, loop: LoopBase) -> DSExperiment:
 
             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