feat: enable to inject diversity cross async multi-trace (#1173)

* init multi_trace_async_ diversity inject

* lint

* add task in context for diversity injection, add abstract class for diversity injection

* lint

* feat: update diversity injection strategy and enhance sibling context handling in experiment generation

* add always inject

---------

Co-authored-by: Xu Yang <peteryang@vip.qq.com>

rdagent/app/data_science/conf.py

@@ -111,6 +111,16 @@ class DataScienceBasePropSetting(KaggleBasePropSetting):
     # inject diverse when start a new sub-trace
     enable_inject_diverse: bool = False
 
+    # inject diverse from other traces when start a new sub-trace
+    enable_cross_trace_diversity: bool = True
+    """Enable cross-trace diversity injection when starting a new sub-trace.
+    This is different from `enable_inject_diverse` which is for non-parallel cases."""
+
+    diversity_injection_strategy: str = (
+        "rdagent.scenarios.data_science.proposal.exp_gen.diversity_strategy.InjectUntilSOTAGainedStrategy"
+    )
+    """The strategy to use for injecting diversity context."""
+
     # enable different version of DSExpGen for multi-trace
     enable_multi_version_exp_gen: bool = False
     exp_gen_version_list: str = "v3,v2"

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

@@ -5,6 +5,7 @@ from rdagent.app.data_science.conf import DS_RD_SETTING
 from rdagent.core.evolving_framework import KnowledgeBase
 from rdagent.core.experiment import Experiment
 from rdagent.core.proposal import ExperimentFeedback, Hypothesis, Trace
+from rdagent.core.utils import import_class
 from rdagent.scenarios.data_science.experiment.experiment import COMPONENT, DSExperiment
 from rdagent.scenarios.data_science.scen import DataScienceScen
 
@@ -60,8 +61,30 @@ class DSTrace(Trace[DataScienceScen, KnowledgeBase]):
 
         self.sota_exp_to_submit: DSExperiment | None = None  # grab the global best exp to submit
 
+        self.uncommitted_experiments: dict[int, DSExperiment] = {}  # loop_id -> DSExperiment
+
+    def should_inject_diversity(self, current_selection: tuple[int, ...] | None = None) -> bool:
+        """
+        Check if diversity context should be injected based on the current selection.
+        This function calls the diversity strategy's should_inject method.
+        """
+        if current_selection is None:
+            current_selection = self.get_current_selection()
+        return (
+            import_class(DS_RD_SETTING.diversity_injection_strategy)().should_inject(self, current_selection)
+            if DS_RD_SETTING.enable_cross_trace_diversity
+            else False
+        )
+
     COMPLETE_ORDER = ("DataLoadSpec", "FeatureEng", "Model", "Ensemble", "Workflow")
 
+    def register_uncommitted_exp(self, exp: DSExperiment, loop_id: int):
+        self.uncommitted_experiments[loop_id] = exp
+
+    def deregister_uncommitted_exp(self, loop_id: int):
+        if loop_id in self.uncommitted_experiments:
+            del self.uncommitted_experiments[loop_id]
+
     def set_sota_exp_to_submit(self, exp: DSExperiment) -> None:
         self.sota_exp_to_submit = exp
 
@@ -88,6 +111,31 @@ class DSTrace(Trace[DataScienceScen, KnowledgeBase]):
         leaves = list(sorted(all_indices - parent_indices))
         return leaves
 
+    def get_sibling_exps(self, current_selection: tuple[int, ...] | None = None):
+        """
+        Get the sibling experiments of the current selection.
+        Include the committed and uncommitted experiments.
+        """
+        if current_selection is None:
+            current_selection = self.get_current_selection()
+        ignore_leaf_idx = [current_selection[0]] if current_selection != self.NEW_ROOT else []
+        sibling_exps = []
+        touched_node_set = set()
+        for idx in range(len(self.dag_parent)):
+            touched_node_set.add(idx)
+            if self.dag_parent[idx] == self.NEW_ROOT:
+                continue
+            for parent in self.dag_parent[idx]:
+                touched_node_set.remove(parent)
+        for loop_idx, exp in self.uncommitted_experiments.items():
+            sibling_exps.append(exp)
+            if (exp_parent_idx := exp.local_selection[0] if exp.local_selection != self.NEW_ROOT else None) is not None:
+                touched_node_set.remove(exp_parent_idx)
+        for idx in touched_node_set:
+            if idx not in ignore_leaf_idx:
+                sibling_exps.append(self.hist[idx][0])
+        return sibling_exps
+
     def sync_dag_parent_and_hist(
         self,
         exp_and_fb: tuple[Experiment, ExperimentFeedback],
@@ -112,6 +160,7 @@ class DSTrace(Trace[DataScienceScen, KnowledgeBase]):
             self.dag_parent.append((current_node_idx,))
         self.hist.append(exp_and_fb)
         self.idx2loop_id[len(self.hist) - 1] = cur_loop_id
+        self.deregister_uncommitted_exp(cur_loop_id)
 
     def retrieve_search_list(
         self,

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

@@ -0,0 +1,68 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from rdagent.scenarios.data_science.proposal.exp_gen.base import DSTrace
+
+
+class DiversityContextStrategy(ABC):
+    """
+    An abstract base class for strategies that determine when to inject
+    cross-trace diversity context into the generation process.
+    """
+
+    @abstractmethod
+    def should_inject(self, trace: DSTrace, local_selection: tuple[int, ...]) -> bool:
+        """
+        Decides whether to inject diversity context based on the current state of the trace
+        and the selection for the next experiment.
+
+        Args:
+            trace: The full DSTrace object.
+            local_selection: The parent node selection for the new experiment.
+
+        Returns:
+            True if context should be injected, False otherwise.
+        """
+        raise NotImplementedError
+
+
+class InjectAtRootStrategy(DiversityContextStrategy):
+    """
+    A strategy that injects diversity context only when creating a new root for a sub-trace.
+    """
+
+    def should_inject(self, trace: DSTrace, local_selection: tuple[int, ...]) -> bool:
+        """Injects only when `local_selection` indicates a new trace root."""
+        return local_selection == trace.NEW_ROOT
+
+
+class InjectUntilSOTAGainedStrategy(DiversityContextStrategy):
+    """
+    A strategy that injects diversity context until the first SOTA (State-of-the-Art)
+    experiment is achieved within the current sub-trace.
+    """
+
+    def should_inject(self, trace: DSTrace, local_selection: tuple[int, ...]) -> bool:
+        """
+        Injects if the sub-trace corresponding to the `local_selection` has not
+        yet produced a successful SOTA experiment.
+        """
+        # If starting a new trace, there's no SOTA yet, so inject.
+        if local_selection == trace.NEW_ROOT:
+            return True
+
+        # Check for SOTA within the specific sub-trace.
+        return trace.sota_experiment(selection=local_selection) is None
+
+
+class AlwaysInjectStrategy(DiversityContextStrategy):
+    """
+    A strategy that always injects diversity context.
+    """
+
+    def should_inject(self, trace: DSTrace, local_selection: tuple[int, ...]) -> bool:
+        """Always returns True to indicate that context should be injected."""
+        return True

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

@@ -16,6 +16,17 @@ scenario_problem:
     {% 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 %}
 
+    {% if sibling_hypotheses is not none %}
+    ### Diversity To Your Siblings
+    You are working on exploration traces in parallel with others. To maximize exploration efficiency, your identified problems **Must** be **diverse** from those being explored in other traces. 
+    Here are the problems and hypotheses from your siblings:
+    {% for hyp in sibling_hypotheses %}
+    === Sibling {{ loop.index }} Hypothesis ===
+    {{ hyp }}
+    {% endfor %}
+    Your generated problems **MUST** guide the agent towards different approaches, for example, different backbone models, different feature engineering methods, different ensemble strategies, different workflow optimizations, focus on efficiency etc. Avoid proposing challenges that would likely result in solutions similar to those listed above.
+    {% 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:
     1. **Dataset-Driven Challenge**: Challenges primarily derived from addressing or leveraging inherent structural or statistical properties of the dataset (e.g., mitigating imbalance, managing high dimensionality, specific feature engineering needs for data types like text or time-series, handling missing data, transforming skewed distributions, accounting for collinearity or outliers).
@@ -91,6 +102,17 @@ feedback_problem:
     4. Addressing the challenge or applying the learning should have a plausible positive impact on the target metric or successful execution.
     5. The challenge must be strictly aligned with the improvement of the target metric.
     
+    {% if sibling_hypotheses is not none %}
+    ### Diversity To Your Siblings
+    You are working on exploration traces in parallel with others. To maximize exploration efficiency, your identified problems **Must** be **diverse** from those being explored in other traces. 
+    Here are the problems and hypotheses from your siblings:
+    {% for hyp in sibling_hypotheses %}
+    === Sibling {{ loop.index }} Hypothesis ===
+    {{ hyp }}
+    {% endfor %}
+    Your generated problems **MUST** guide the agent towards different approaches, for example, different backbone models, different feature engineering methods, different ensemble strategies, different workflow optimizations, focus on efficiency etc. Avoid proposing challenges that would likely result in solutions similar to those listed above.
+    {% endif %}
+    
     {% if problem_output_format is not none %}
     ### Output Format
     {{ problem_output_format }}
@@ -193,6 +215,17 @@ hypothesis_gen:
       - **Risk-Reward Balance (Score: 1-10):** Considering the potential for significant improvement (reward) versus the probability of failure, negative side-effects, or excessive resource consumption (risk), how optimal is this balance? A high score indicates a favorable balance.
       - **Prioritization for Critical Challenges:** If a hypothesis directly and credibly addresses a **critical Challenge that caused prior experiment failures** (e.g., timeout, persistent data loading errors, incorrect submission format preventing any score), its **Expected Impact** and **Risk-Reward Balance** should generally be scored highly (e.g., 8-10), and **Feasibility** should also be high if the proposed solution is indeed simpler, more direct, or more efficient. This ensures such critical hypotheses are prioritized.
 
+    {% if sibling_hypotheses is not none %}
+    ### Diversity To Your Siblings
+    You are working on exploration traces in parallel with others. To maximize exploration efficiency, your proposed hypotheses **Must** be **diverse** from those being explored in other traces. 
+    Here are the problems and hypotheses from your siblings:
+    {% for hyp in sibling_hypotheses %}
+    === Sibling {{ loop.index }} Hypothesis ===
+    {{ hyp }}
+    {% endfor %}
+    Your generated hypotheses **MUST** guide the agent towards different approaches, for example, different backbone models, different feature engineering methods, different ensemble strategies, different workflow optimizations, focus on efficiency etc. Avoid proposing hypotheses that are similar to those listed above.
+    {% endif %}
+
     {% 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.
@@ -364,6 +397,16 @@ hypothesis_rewrite:
       - Apply critique insights to enhance sound innovative ideas while avoiding repeated fundamental failures identified in the analysis.
       - **Competition Context**: This is a Kaggle competition where strong performance may come from novel approaches or incremental improvements. Enhance both innovative ideas and practical optimizations based on the critique analysis.
     
+    {% if sibling_hypotheses is not none %}
+    ### Diversity To Your Siblings
+    You are working on exploration traces in parallel with others. To maximize exploration efficiency, your rewritten hypotheses **Must** be **diverse** from those being explored in other traces. 
+    Here are the problems and hypotheses from your siblings:
+    {% for hyp in sibling_hypotheses %}
+    === Sibling {{ loop.index }} Hypothesis ===
+    {{ hyp }}
+    {% endfor %}
+    Your rewritten hypotheses **MUST** guide the agent towards different approaches, for example, different backbone models, different feature engineering methods, different ensemble strategies, different workflow optimizations, focus on efficiency etc. Avoid proposing hypotheses that are similar to those listed above.
+    {% endif %}
     
     {% if rewrite_output_format is not none %}
     ## Output Format
@@ -406,6 +449,17 @@ task_gen:
 
     Your primary goal is to generate a detailed, step-by-step **sketch or refinement plan** for a new data processing and modeling pipeline, specifically for the main workflow script (`main.py`), that effectively implements the `Proposed Hypothesis`. This sketch will guide a developer to write the code correctly.
 
+    {% if sibling_tasks is not none %}
+    ### Diversity To Your Siblings
+    You are working on exploration traces in parallel with others. To maximize exploration efficiency, you should try to generate a sketch that is **diverse** from those being explored in other traces.
+    Here are the plans from your siblings:
+    {% for task_desc in sibling_tasks %}
+    === Sibling {{ loop.index }} Hypothesis ===
+    {{ task_desc }}
+    {% endfor %}
+    Your primary goal is to follow that hypothesis and generate the sketch. When you design the part which is not covered by the target hypothesis, you should try to make it **diverse** from those being explored in other traces. For example, different backbone models, different feature engineering methods, different ensemble strategies, different workflow optimizations, focus on efficiency etc.
+    {% endif %}
+
     # BACKGROUND CONTEXT: Pipeline Implementation Standards & Constraints
 
     The `main.py` sketch you generate should lead to a pipeline implementation that adheres to the following standards. These are guiding principles for the final *outcome* of your sketch:

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

@@ -485,12 +485,15 @@ class DSProposalV2ExpGen(ExpGen):
         scenario_desc: str,
         sota_exp_desc: str,
         exp_gen_plan: Dict,
+        sibling_exp: List[DSExperiment] | None = None,
     ) -> Dict:
+        sibling_hypotheses = [exp.hypothesis for exp in sibling_exp] if sibling_exp else None
         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,
+            sibling_hypotheses=sibling_hypotheses,
         )
         user_prompt = T(".prompts_v2:scenario_problem.user").r(
             scenario_desc=scenario_desc,
@@ -513,13 +516,20 @@ class DSProposalV2ExpGen(ExpGen):
         return problems
 
     def identify_feedback_problem(
-        self, scenario_desc: str, exp_feedback_list_desc: str, sota_exp_desc: str, inject_diverse: bool = False
+        self,
+        scenario_desc: str,
+        exp_feedback_list_desc: str,
+        sota_exp_desc: str,
+        inject_diverse: bool = False,
+        sibling_exp: List[DSExperiment] | None = None,
     ) -> Dict:
+        sibling_hypotheses = [exp.hypothesis for exp in sibling_exp] if sibling_exp else None
         sys_prompt = T(".prompts_v2:feedback_problem.system").r(
             problem_output_format=(
                 T(".prompts_v2:output_format.problem").r() if not self.supports_response_schema else None
             ),
             inject_diverse=inject_diverse,
+            sibling_hypotheses=sibling_hypotheses,
         )
         user_prompt = T(".prompts_v2:feedback_problem.user").r(
             scenario_desc=scenario_desc,
@@ -550,6 +560,7 @@ class DSProposalV2ExpGen(ExpGen):
         exp_feedback_list_desc,
         inject_diverse,
         exp_gen_plan,
+        sibling_exp: List[DSExperiment] | None = None,
     ) -> 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
@@ -562,6 +573,7 @@ class DSProposalV2ExpGen(ExpGen):
                 scenario_desc=scenario_desc,
                 sota_exp_desc=sota_exp_desc,
                 exp_gen_plan=exp_gen_plan,
+                sibling_exp=sibling_exp,
             )
             for problem_name in scen_problems:
                 scen_problems[problem_name]["label"] = "SCENARIO_PROBLEM"
@@ -592,6 +604,7 @@ class DSProposalV2ExpGen(ExpGen):
         inject_diverse: bool = False,
         exp_gen_plan: Optional[Dict] = None,
         packages_prompt: str = "",
+        sibling_exp: List[DSExperiment] | None = None,
     ) -> Dict:
         problem_formatted_str = ""
         for i, (problem_name, problem_dict) in enumerate(problems.items()):
@@ -602,6 +615,11 @@ class DSProposalV2ExpGen(ExpGen):
                 idea_formatted_str = DSIdea(problem_dict["idea"]).to_formatted_str()
                 problem_formatted_str += f"Sampled Idea by user: \n{idea_formatted_str}\n"
             problem_formatted_str += "\n\n"
+        sibling_hypotheses = [exp.hypothesis for exp in sibling_exp] if sibling_exp else None
+
+        # add available packages prompt
+        if packages_prompt:
+            problem_formatted_str += f"\n{packages_prompt}\n"
 
         # add available packages prompt
         if packages_prompt:
@@ -617,6 +635,7 @@ class DSProposalV2ExpGen(ExpGen):
             enable_idea_pool=enable_idea_pool,
             inject_diverse=inject_diverse,
             plan=exp_gen_plan,
+            sibling_hypotheses=sibling_hypotheses,
         )
         user_prompt = T(".prompts_v2:hypothesis_gen.user").r(
             scenario_desc=scenario_desc,
@@ -741,11 +760,14 @@ class DSProposalV2ExpGen(ExpGen):
         sota_exp_desc: str,
         exp_feedback_list_desc: str,
         packages_prompt: str = "",
+        sibling_exp: List[DSExperiment] | None = None,
     ) -> Dict:
         """
         Generate improved hypotheses based on critique feedback for each original hypothesis.
         Returns a dict with the same keys as hypothesis_dict, containing improved versions.
         """
+        sibling_hypotheses = [exp.hypothesis for exp in sibling_exp] if sibling_exp else None
+
         hypothesis_critique_pairs = ""
         for i, problem_name in enumerate(hypothesis_dict.keys()):
             hypothesis_data = hypothesis_dict[problem_name]
@@ -772,6 +794,7 @@ class DSProposalV2ExpGen(ExpGen):
                 enable_scale_check=DS_RD_SETTING.enable_scale_check
             ),
             enable_scale_check=DS_RD_SETTING.enable_scale_check,
+            sibling_hypotheses=sibling_hypotheses,
         )
         user_prompt = T(".prompts_v2:hypothesis_rewrite.user").r(
             scenario_desc=scenario_desc,
@@ -915,6 +938,7 @@ class DSProposalV2ExpGen(ExpGen):
         pipeline: bool,
         failed_exp_feedback_list_desc: str,
         fb_to_sota_exp: ExperimentFeedback | None = None,
+        sibling_exp: List[DSExperiment] | None = None,
     ) -> DSExperiment:
         if pipeline:
             component_info = get_component("Pipeline")
@@ -922,11 +946,14 @@ class DSProposalV2ExpGen(ExpGen):
             component_info = get_component(hypotheses[0].component)
         data_folder_info = self.scen.processed_data_folder_description
         workflow_check = not pipeline and hypotheses[0].component != "Workflow"
+
+        sibling_tasks = [exp.pending_tasks_list[0][0].description for exp in sibling_exp] if sibling_exp else []
         sys_prompt = T(".prompts_v2:task_gen.system").r(
             task_output_format=component_info["task_output_format"] if not self.supports_response_schema else None,
             component_desc=component_desc,
             workflow_check=workflow_check,
             metric_name=self.scen.metric_name,
+            sibling_tasks=sibling_tasks,
         )
         user_prompt = T(".prompts_v2:task_gen.user").r(
             scenario_desc=scenario_desc,
@@ -1070,6 +1097,8 @@ class DSProposalV2ExpGen(ExpGen):
         # add available packages prompt
         packages_prompt = get_available_packages_prompt()
 
+        sibling_exp = trace.get_sibling_exps() if trace.should_inject_diversity() else None
+
         # Step 1: Identify problems
         all_problems = self.identify_problem(
             current_sub_trace=trace.get_parent_exps(),
@@ -1078,6 +1107,7 @@ class DSProposalV2ExpGen(ExpGen):
             exp_feedback_list_desc=exp_feedback_list_desc,
             inject_diverse=inject_diverse,
             exp_gen_plan=plan.get("exp_gen") if plan else None,
+            sibling_exp=sibling_exp,
         )
 
         # Step 1.5: Sample ideas from idea pool
@@ -1102,6 +1132,7 @@ class DSProposalV2ExpGen(ExpGen):
             inject_diverse=inject_diverse,
             exp_gen_plan=plan.get("exp_gen") if plan else None,
             packages_prompt=packages_prompt,
+            sibling_exp=sibling_exp,
         )
         if not pipeline:
             sota_exp_model_file_count = len(
@@ -1146,6 +1177,7 @@ class DSProposalV2ExpGen(ExpGen):
                     sota_exp_desc=sota_exp_desc,
                     exp_feedback_list_desc=exp_feedback_list_desc,
                     packages_prompt=packages_prompt,
+                    sibling_exp=sibling_exp,
                 )
                 logger.info(f"Successfully completed hypothesis critique and rewrite process")
             except Exception as e:
@@ -1178,4 +1210,5 @@ class DSProposalV2ExpGen(ExpGen):
             pipeline=pipeline,
             failed_exp_feedback_list_desc=failed_exp_feedback_list_desc,
             fb_to_sota_exp=fb_to_sota_exp,
+            sibling_exp=sibling_exp,
         )

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

@@ -130,6 +130,9 @@ class ParallelMultiTraceExpGen(ExpGen):
                 )
                 exp.set_local_selection(local_selection)
                 exp.plan = ds_plan
+
+                # Register the newly created experiment before returning
+                trace.register_uncommitted_exp(exp, loop.loop_idx)
                 return exp
 
             await asyncio.sleep(1)