feat: enable LLM‑based hypothesis selection with time‑aware prompt & colored logging (#1122)

* add hypo select by llm (without time) * add time_info and log color * select no smooth * 2 hypo * change select * small change * fix bug * fix bug and add hypothesis router and begin flag * fix bug v1 * fix bug v2 * fix feedback * add new model * add filter * fix bug v3 * fix bug v4 * change prompts v2 * fix bug v5 * fix bug v6 * fix hypo * fix some bug(sota socre, prompts, ensemble prompts ) and add path legth. * fix bug v7 * fix bug v8 * fix bug v9 * fix bug v10 * reset to v10 and refine * fix: translate to english * fix bug * fix: use differnet increase_stage for coder & runner * fix: use different timeout_increase_stage for coder and runner. * fix: revert logger * fix: remove duplicate content * feat: implement LLM-driven extra hypothesis selection and adjust logic * refactor: relocate Hypothesis models below TraceChallenges * remove torch fix bug * fix small bug * refactor: prefix internal methods with underscore for llm-based hypothesis selection * add fix_seed_and_data_split and enable_simple_hypothesis * lint * refine proposal config order * code review; comments * lint * merge int to float --------- Co-authored-by: jingyuanlm <842442862@qq.com> Co-authored-by: Young <afe.young@gmail.com>
2026-07-27 23:47:46 +00:00 · 2025-08-27 18:36:51 +08:00
parent 3df94e920c
commit 90dd2f7b9b
7 changed files with 593 additions and 42 deletions
@@ -42,6 +42,11 @@ class DataScienceBasePropSetting(KaggleBasePropSetting):
    full_timeout: int = 3600
    """The timeout limit for running on full data"""

+    #### model dump
+    enable_model_dump: bool = False
+    enable_doc_dev: bool = False
+    model_dump_check_level: Literal["medium", "high"] = "medium"
+
    ### specific feature

    ### notebook integration
@@ -63,11 +68,6 @@ class DataScienceBasePropSetting(KaggleBasePropSetting):
    use_raw_description: bool = False
    show_nan_columns: bool = False

-    #### model dump
-    enable_model_dump: bool = False
-    enable_doc_dev: bool = False
-    model_dump_check_level: Literal["medium", "high"] = "medium"
-
    ### knowledge base
    enable_knowledge_base: bool = False
    knowledge_base_version: str = "v1"
@@ -126,7 +126,7 @@ class DataScienceBasePropSetting(KaggleBasePropSetting):
    exp_gen_version_list: str = "v3,v2"

    #### multi-trace: time for final multi-trace merge
-    merge_hours: int = 0
+    merge_hours: float = 0
    """The time for merge"""

    #### multi-trace: max SOTA-retrieved number, used in AutoSOTAexpSelector
@@ -140,19 +140,44 @@ class DataScienceBasePropSetting(KaggleBasePropSetting):

    model_architecture_suggestion_time_percent: float = 0.75
    allow_longer_timeout: bool = False
-    longer_timeout_by_llm: bool = False
+    coder_enable_llm_decide_longer_timeout: bool = False
+    runner_enable_llm_decide_longer_timeout: bool = False
    coder_longer_timeout_multiplier_upper: int = 3
    runner_longer_timeout_multiplier_upper: int = 2
-    timeout_increase_stage: float = 0.3
+    coder_timeout_increase_stage: float = 0.3
+    runner_timeout_increase_stage: float = 0.15
    show_hard_limit: bool = True

-    #### hypothesis critique and rewrite
-    enable_hypo_critique_rewrite: bool = True
-    """Enable hypothesis critique and rewrite stages for improving hypothesis quality"""
-    enable_scale_check: bool = False
-
    #### enable runner code change summary
    runner_enable_code_change_summary: bool = True

+    ### Proposal workflow related
+
+    #### Hypothesis Generate related
+    enable_simple_hypothesis: bool = False
+    """If true, generate simple hypothesis, no more than 2 sentences each."""
+
+    enable_generate_unique_hypothesis: bool = False
+    """Enable generate unique hypothesis. If True, generate unique hypothesis for each component. If False, generate unique hypothesis for each component."""
+
+    #### hypothesis critique and rewrite
+    enable_hypo_critique_rewrite: bool = False
+    """Enable hypothesis critique and rewrite stages for improving hypothesis quality"""
+    enable_scale_check: bool = False
+
+    ##### select related
+    ratio_merge_or_ensemble: int = 70
+    """The ratio of merge or ensemble to be considered as a valid solution"""
+    llm_select_hypothesis: bool = False
+    """Whether to use LLM to select hypothesis. If True, use LLM selection; if False, use the existing ranking method."""
+
+    #### Task Generate related
+    fix_seed_and_data_split: bool = False
+

 DS_RD_SETTING = DataScienceBasePropSetting()
+
+# enable_cross_trace_diversity 和 llm_select_hypothesis should not be true at the same time
+assert not (
+    DS_RD_SETTING.enable_cross_trace_diversity and DS_RD_SETTING.llm_select_hypothesis
+), "enable_cross_trace_diversity and llm_select_hypothesis cannot be true at the same time"
@@ -31,9 +31,13 @@ exp_feedback:
      - If overfitting is detected, provide a detailed analysis explaining how and why it occurs, referencing scenario description, code implementation, and validation scores to support your findings.
    - If such discrepancies or risks are found:
      - Clearly document these issues in `Reasoning`, referencing both scenario description and code implementation—not just validation scores.
-      - Set `"Evaluation Aligned With Task": "no"` and `"Replace Best Result": "no"`.
-      - Begin your `reasoning` with `[Evaluation error]`, explicitly stating the evaluation alignment issues causing experiment failure.
-    - If evaluation alignment passes, set `"Evaluation Aligned With Task": "yes"`, and then proceed to Step 3.
+        - Severity-based handling:
+         - Severe risk — likely to invert or invalidate the performance trend between validation and test (e.g., strong overfitting, label leakage, test distribution shift):
+           - Set "Evaluation Aligned With Task": "no" and "Replace Best Result": "no".
+           - Begin your reasoning with [Evaluation error], explicitly stating the evaluation alignment issues causing experiment failure.
+         - Mild/moderate risk — may cause slightly optimistic or biased validation scores but is unlikely to change the relative performance trend (e.g., scaling or PCA fit on full training data that’s also applied consistently to test):
+          - Set "Evaluation Aligned With Task": "yes" but note the potential bias in Reasoning.
+           - Proceed to Step 3 for result comparison.

    Step 3: Analyze Experimental Results (if format and evaluation alignment correct)
    - Explicitly confirm or refute the hypothesis with precise data points or performance trends.
@@ -227,7 +227,12 @@ class DSTrace(Trace[DataScienceScen, KnowledgeBase]):
    ) -> list[tuple[DSExperiment, ExperimentFeedback]]:
        """
        Retrieve a list of experiments and feedbacks based on the return_type.
+
+        return_type:
+            - "sota": experiments that have true decision feedback
        """
+        # TODO: SOTA is a ver confusing name
+
        search_list = self.retrieve_search_list(search_type, selection=selection)
        final_component = self.COMPLETE_ORDER[-1]
        has_final_component = True if DS_RD_SETTING.coder_on_whole_pipeline else False
@@ -235,7 +240,8 @@ class DSTrace(Trace[DataScienceScen, KnowledgeBase]):
        failed_exp_and_feedback_list_after_sota = []
        for exp, fb in search_list:
            if has_final_component:
-                if fb.decision:
+                # FIXME: fb should not be None, but there is a potential bug in the code.
+                if getattr(fb, "decision", False):
                    SOTA_exp_and_feedback_list.append((exp, fb))
                    failed_exp_and_feedback_list_after_sota = []
                else:
@@ -214,7 +214,25 @@ hypothesis_gen:
      - **Feasibility (Score: 1-10):** How easily and practically can this hypothesis be implemented and *run to completion* within the existing SOTA codebase and operational constraints (e.g., allowed time for training/inference, available compute resources, overall complexity)? Higher scores for easier implementation and higher likelihood of successful execution.
      - **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 enable_simple_hypothesis%}
+    3. Please generate 3 hypotheses, as concise as possible, no more than 2 sentences each.
+    {% endif %}
+    {%if generate_unique_hypothesis %}
+    We are now at the beginning stage. Please generate hypotheses that are as unique as possible.
+    Each hypothesis should handle a different component. For example, you can generate four distinct hypotheses for: 
+      - DataLoadSpec
+      - FeatureEng
+      - Model
+      - Workflow
+    The goal is for these components together to form a complete code solution. Avoid generating complex ensemble methods (e.g., 5-fold CV or stacked models) at this stage.  
+    Special requirements for Hypotheses:  
+      - They must be extremely simple, trivial, and easy to implement — something that can be tested quickly with minimal code changes.  
+      - Avoid "trick-like" operations, such as freezing layers in the model.  
+    - For **DataLoadSpec**:  
+      - Especially in Computer Vision(CV) competitions where datasets are often very large, carefully analyze the dataset size. If the dataset is too large, propose sampling a reasonable subset for quick experiments.  
+      - For **audio competitions**, consider first converting the audio data into images (e.g., spectrograms) and then applying CV-based methods for modeling.
+    {% 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 proposed hypotheses **Must** be **diverse** from those being explored in other traces. 
@@ -435,6 +453,188 @@ hypothesis_rewrite:
    {{ packages_prompt }}
    {% endif %}

+hypothesis_select:
+  system: |-
+    You are a Kaggle Grandmaster with deep expertise in model evaluation and decision making.  
+    Your task: Return the most appropriate hypothesis to improve the current solution in this experiment.
+    ## Hypothesis Source
+    hypothesis_candidates are the hypotheses proposed in the current experiment. Please give them priority:
+    {{hypothesis_candidates}}
+
+    {%if sota_flag %}
+    SOTA score: {{current_sota_score}}
+    {% if current_sota_score_in_current_trace == -1 %}
+    Current SOTA score in this experiment: None.
+    {% else %}
+    Current SOTA score in this experiment: {{ current_sota_score_in_current_trace }}
+    {% endif %}
+
+    {% if selected_extra_hypo_l and selected_extra_hypo_l|length > 0 %}
+    The following are additional hypotheses that have been approved by other experiments. 
+    If any of these hypotheses have a SOTA score significantly higher than the current SOTA score in this experiment, you may want to prioritize considering them:  
+    Additional hypotheses (may include those corresponding to the SOTA score):
+    {% for i, item in enumerate(selected_extra_hypo_l) %}
+    {{ i+1 }}. {{ item[0] }} (item[1]: {{ "%.3f"|format(item[1]) }})
+    {% endfor %}
+    {% endif %}
+    {% else %}
+
+    {% if current_sota_score_in_current_trace == -1 %}
+    {% if selected_extra_hypo_l and selected_extra_hypo_l|length > 0 %}
+    The current SOTA score in this experiment is unavailable. Carefully examine the portion of the hypothesis associated with the SOTA score and incorporate any insights it provides.
+    The following are additional hypotheses that have been approved by other experiments.  
+    You can also serve as references and are part of the Hypothesis Source to help you quickly reach or surpass the SOTA score:  
+    Additional hypotheses (may include those corresponding to the SOTA score):
+    {% for i, item in enumerate(selected_extra_hypo_l) %}
+    {{ i+1 }}. {{ item[0] }} (item[1]: {{ "%.3f"|format(item[1]) }})
+    {% endfor %}
+    {% endif %}
+    {% endif %}
+
+    {%endif %}
+
+
+    - The list `hypothesis_candidates` is for REFERENCE ONLY.
+    - You may:
+      1. Select one hypothesis directly from the candidates.
+      2. Modify an existing hypothesis from the candidates.
+      3. Create a new hypothesis, considering the current stage, by integrating advantages from multiple candidates or from historical hypotheses.
+
+    ## Hypothesis Generation Guidelines
+    ## Stage Constraints
+
+    {% if use_ratio < 10 %}
+    ### Stage = Draft
+    - This stage is focused on rapid, easy-to-implement hypotheses. Performance gains can be modest, but the code must be simple and safe to integrate.
+    - You may take one of three actions:
+
+      1. **Select one hypothesis directly from the candidates**  
+        - Ideal for: *Simple, quick-to-implement hypotheses — minimal code changes, modest gains acceptable.*  
+        - Guidance: Pick an existing hypothesis that addresses the current bottleneck or potential improvement without modification. This is the fastest way to produce working code.
+
+      2. **Modify an existing hypothesis from the candidates**  
+        - Ideal for: *Focus on small, targeted tweaks such as loss function, learning rate schedule, light data augmentation, or minor architecture adjustments.*  
+        - Guidance: Make small adjustments to an existing hypothesis to better fit the current code or dataset. Examples include:  
+          - Tuning hyperparameters (including learning rate, batch size, and number of epochs)  
+          - Adjusting the loss function  
+          - Applying lightweight augmentations  
+          - Minor architecture modifications  
+
+          **High-priority suggestions based on competition type:**  
+          - **CV competitions:** Consider using larger image sizes and the latest model architectures (e.g., Swin Transformer, Vision Transformer (ViT), EfficientNetV2).  
+          - **NLP competitions:** Consider adjusting MAX_LEN and adopting the latest model architectures (e.g., DeBERTa v3-large, RoBERTa).  
+          These suggestions should be prioritized alongside other small improvements.
+
+      3. **Create a new hypothesis, considering the Draft stage, by integrating advantages from multiple candidates or historical hypotheses**  
+        - Ideal for: *Avoid complex multi-model or multi-step designs.*  
+        - Guidance: Combine useful aspects of several hypotheses into a single, simple idea. Ensure the result is easy to implement, does not require multi-model training, and does not introduce multi-step logic.
+    {% elif use_ratio > ratio_merge_or_ensemble %}
+
+
+    ### Stage = Ensemble
+    - This stage focuses on maximizing overall performance by combining multiple models or hypotheses. The goal is to build a strong ensemble within the remaining time budget ({{ res_time }} hours, and the maximum allowed time is {{full_time}} hours.). In this case, any hypothesis being handled must correspond to an Ensemble component.
+    - **Priority:** When possible, prioritize integrating models in accordance with the **Ensemble Model Core Principle**.
+
+    {%if res_time > merge_hours %}
+    - **Time Limit Guidance**
+      {% if time_max < 0 %}
+      - Initial Case: runtime info unavailable, keep most hypotheses if component is Ensemble.
+      - Remove only those clearly excessive (e.g., > {{ full_time }} hours) or overly complex.
+      {% elif time_max >= full_time * 0.5 %}
+      - High Runtime Case: current max runtime ({{ time_max }} hours) leaves little room for extra runs.
+      - Avoid high-fold or heavy ensembles.
+      - Maximum recommended folds: {{ (full_time // time_max) | int }}
+      - Remove hypotheses clearly excessive (> {{ full_time }} hours)
+      {% else %}
+      - Low Runtime Case: current max runtime ({{ time_max }} hours) is far from the time limit.
+      - Prefer hypotheses with runtimes ≤ {{ full_time }} hours.
+      - Hypotheses slightly above {{ time_max }} hours can be retained only with strong justification.
+      {% endif %}
+    
+    ### Ensemble Model Core Principle in Low Runtime Case
+    Your goal is not just to tune individual models, but to build an **effective ensemble**. Make design decisions that lead to **strong overall ensemble performance**, not just strong base models.
+    Please note: you are operating under a time budget dedicated to ensemble training of {{res_time}} seconds, and the maximum allowed time is {{full_time}} seconds.
+  
+    Please take the remaining {{res_time}} seconds to carefully consider and design the most reasonable and optimal ensemble models based on your current progress.
+    Assume training a single model takes about 1 hour. For example, if you have roughly twice that time left, you can try training multiple models with different random seeds or data splits to reuse time effectively.
+    If you have more time, you might consider training a multi-fold ensemble. Use your judgment to decide how many folds or seeds fit within your remaining time budget.
+
+    ### 2. Training-Time Resource Allocation
+    - You may use **multiple folds** if justified, but you must **ensure the full pipeline completes within runtime limits**.
+    - Avoid reducing base model quality just to save time. For example:
+      - Freezing large parts of the model (e.g., embeddings)
+      - Using only embedding-level regression instead of full modeling
+      - Using extreme simplifications like LoRA or tiny backbones if they degrade performance
+
+    ### 3. Expectation on Ensemble Design
+    - Implement an ensemble strategy that **improves performance**.
+      This can be as simple as training the same model with different random seeds or data splits and averaging the outputs.
+      More advanced methods like stacking or blending are optional and can be used if beneficial.
+      Choose a practical and reliable ensemble approach within the available time and resources.
+    - Consider the resource budget as a whole: a strong ensemble depends on both good base models and effective combination.
+
+    ### 4. Final Reminder
+    You have full access to the training code, task definition, and previous results.
+    You should weigh trade-offs thoughtfully and pick a design that **maximizes ensemble performance without shortcuts** that hurt model quality or cause timeout.
+    - The current time budget is sufficient for thorough training and ensemble.
+    - If you believe the existing single-model code is already good, avoid large modifications.
+    - Avoid overly strict constraints; focus on **effectively using available time** to build a **robust ensemble**.
+
+    {% endif %}
+
+    According to the previous Time Limit Guidance. You may take one of three actions, considering the remaining time and runtime guidance:
+
+      1. **Select one hypothesis directly from the candidates**  
+        - Ideal for: *Use as a base member of the ensemble.*  
+        - Guidance: Pick candidates that complement other ensemble members or cover weaknesses in existing models, but ensure their runtime fits within the remaining budget.
+
+      2. **Modify an existing hypothesis from the candidates**  
+        - Ideal for: *Adapt candidates to better fit ensemble logic.*  
+        - Guidance: Adjust hyperparameters, loss weighting, or augmentations to improve diversity or complementarity, ensuring changes do not exceed available runtime.
+
+      3. **Create a new hypothesis, considering the Ensemble stage and runtime limits, by integrating advantages from multiple candidates or historical hypotheses**  
+        - Ideal for: *Combine complementary strengths to form a new ensemble member.*  
+        - Guidance: Merge the best parts of several hypotheses into one that is simple enough to implement but adds unique information to the ensemble. Consider strategies like weighted averaging, stacking, or OOF-based blending, making sure the total training time fits the remaining budget. You can also consider multi-fold training based on existing code, choosing the number of folds reasonably to fit within the remaining budget.
+
+    {% else %}
+
+    ### Stage = Improvement
+    - This stage focuses on achieving meaningful improvement without overcomplicating code. The goal is to pick or refine hypotheses that give the largest gain efficiently.
+
+    - You may take one of three actions:
+
+      1. **Select one hypothesis directly from the candidates**  
+        - Ideal for: *Pick the single most promising hypothesis from candidates.*  
+        - Guidance: Choose the hypothesis with the highest expected impact. Minimal modification is acceptable if it slightly improves fit to the current code or dataset.
+
+      2. **Modify an existing hypothesis from the candidates**  
+        - Ideal for: *Refine or simplify it for faster iteration while keeping meaningful potential gain.*  
+        - Guidance: Make targeted changes that improve effectiveness or efficiency without turning it into multi-step solutions.
+              Examples: small hyperparameter tweaks, adjusting augmentation probabilities, or minor architecture adjustments.
+              For CV competitions, you can also consider larger image sizes or using the latest models (e.g., Swin Transformer, Vision Transformer (ViT), EfficientNetV2).
+              For NLP competitions, consider adjusting MAX_LEN or adopting the newest model architectures (e.g., DeBERTa v3-large, RoBERTa).
+      3. **Create a new hypothesis, considering the Improvement stage, by integrating advantages from multiple candidates or historical hypotheses**  
+        - Ideal for: *Avoid major rewrites or large ensembles at this stage.*  
+        - Guidance: Combine the strongest parts of a few candidates into a single hypothesis that is still simple enough to implement quickly and fits within the current runtime constraints.
+
+    {% endif %}
+
+
+    {% if hypothesis_output_format is not none %}
+    ## Final Output Format in JSON Schema:
+    {{ hypothesis_output_format }}
+    {% endif %}
+    
+
+  user: |-
+    # Scenario Description
+    {{ scenario_desc }}
+
+    # Previous Experiments and Feedbacks
+    {{ exp_and_feedback_list_desc }}
+
+    # Current SOTA Implementation
+    {{ sota_exp_desc }}

 task_gen:
  system: |-
@@ -497,7 +697,24 @@ task_gen:
    8. **Submission File (`submission.csv`)**: Generate `submission.csv` in the **exact format** required (column names, order, data types), as detailed in the '====== Submission Format ======' section of the Competition Scenario Description (DO NOT read the sample_submission.csv file directly in the code). This is a critical step.
    9. **Preferred Packages Notes**:
      {% include "scenarios.data_science.share:guidelines.package_selection" %}
-    
+
+    {%if fix_seed_and_data_split %}
+    10. File Handling & DataFrame Generation: Generate a pandas DataFrame with columns [“id”, “path”, “fold”].
+      - id: a unique identifier for each sample.
+      - path: the file path of the corresponding sample.
+      - split: indicates the assignment of each sample for data splitting. Two modes are supported:
+        - K-Fold (optional): assign integers 0, 1, …, K-1 for each fold.
+        - Train/Test Split (optional): assign "train" or "test" for each sample according to the split ratio (e.g., 8:2).
+        - Ensure reproducibility: the DataFrame must be generated exactly the same way every time the script runs, e.g., by fixing the random seed 42.
+      Data Splitting: use this DataFrame to perform dataset splitting, selecting samples for training and testing based on the fold column.
+    11. Random Seed for Model Training:
+      - If training neural networks, ensure the initial weights and all random operations use a fixed seed of 42 (e.g., torch.manual_seed(42), numpy.random.seed(42), random.seed(42)).
+      - If training machine learning models such as LightGBM, XGBoost, or scikit-learn estimators, absolutely ensure the random seed is fixed (e.g., `random_state=42`) to guarantee reproducibility.
+      - This is mandatory: all aspects of the experiment must be fully reproducible and aligned, including dataset splits and random seeds;
+      - For multi-fold training, use out-of-fold (OOF) predictions as validation scores and save them as an oof file.
+    12. Hypothesis Handling: At the initial stage, multiple hypotheses may be proposed simultaneously. If some hypotheses overlap, select the most promising one for implementation and ignore redundant overlapping hypotheses. Each implemented hypothesis should remain an independent task.
+    Ensure reproducibility: the DataFrame must be generated exactly the same way every time the script runs, regardless of system or runtime conditions (e.g., by fixing the random seed).
+    {% endif %}
    ## Package Declaration
    At the end of your design, **you MUST** provide a key `packages` in the final JSON output.  
    It should be an **array of PyPI package names** (strings) that you expect to `import` in the forthcoming implementation.  
@@ -729,4 +946,10 @@ output_format:
      }
    }

+  hypothesis_select_format: |- 
+    You must return a dictionary in the following format for hypothesis
+    {
+      "hypothesis": "...",  
+      "component": "..."  // Must be one of: 'DataLoadSpec', 'FeatureEng', 'Model', 'Workflow', 'Ensemble'
+    }

@@ -1,7 +1,10 @@
 import json
+import math
+from datetime import timedelta
 from enum import Enum
 from typing import Any, Dict, List, Optional, Tuple

+import numpy as np
 import pandas as pd
 from pydantic import BaseModel, Field

@@ -15,6 +18,7 @@ from rdagent.components.coder.data_science.workflow.exp import WorkflowTask
 from rdagent.core.proposal import ExpGen
 from rdagent.core.scenario import Scenario
 from rdagent.log import rdagent_logger as logger
+from rdagent.log.timer import RD_Agent_TIMER_wrapper
 from rdagent.oai.llm_utils import APIBackend, md5_hash
 from rdagent.scenarios.data_science.dev.feedback import ExperimentFeedback
 from rdagent.scenarios.data_science.experiment.experiment import DSExperiment
@@ -27,10 +31,14 @@ from rdagent.scenarios.data_science.proposal.exp_gen.planner import (
    DSExperimentPlan,
    RD_Agent_TIMER_wrapper,
 )
+from rdagent.scenarios.data_science.proposal.exp_gen.select.submit import (
+    BestValidSelector,
+)
 from rdagent.scenarios.data_science.proposal.exp_gen.utils import (
    get_available_packages_prompt,
    get_packages,
 )
+from rdagent.scenarios.kaggle.kaggle_crawler import get_metric_direction
 from rdagent.utils.agent.tpl import T
 from rdagent.utils.repo.diff import generate_diff_from_dict
 from rdagent.utils.workflow import wait_retry
@@ -250,6 +258,13 @@ class HypothesisDetail(BaseModel):
    evaluation: HypothesisEvaluation = Field(description="Evaluate the quality of the hypothesis.")


+class HypothesisSimple(BaseModel):
+    hypothesis: str = Field(
+        description="The statement of the hypothesis. It could be a design of a new component, or a concise, testable statement derived from previous experimental outcomes."
+    )
+    component: HypothesisComponent = Field(description="The component tag of the hypothesis.")
+
+
 class HypothesisList(BaseModel):
    deduplicated_challenges: List[str] = Field(
        description="A list of deduplicated challenge captions. Each must retain its original wording. If multiple captions are semantically identical, keep the first one."
@@ -601,6 +616,7 @@ class DSProposalV2ExpGen(ExpGen):
        problems: dict,
        pipeline: bool,
        enable_idea_pool: bool,
+        is_new_tree: bool,
        inject_diverse: bool = False,
        exp_gen_plan: Optional[Dict] = None,
        packages_prompt: str = "",
@@ -635,6 +651,8 @@ class DSProposalV2ExpGen(ExpGen):
            enable_idea_pool=enable_idea_pool,
            inject_diverse=inject_diverse,
            plan=exp_gen_plan,
+            generate_unique_hypothesis=DS_RD_SETTING.enable_generate_unique_hypothesis and is_new_tree,
+            enable_simple_hypothesis=DS_RD_SETTING.enable_simple_hypothesis,
            sibling_hypotheses=sibling_hypotheses,
        )
        user_prompt = T(".prompts_v2:hypothesis_gen.user").r(
@@ -902,6 +920,255 @@ class DSProposalV2ExpGen(ExpGen):
        )
        return index_to_pick_pool_list[reproducible_int]

+    # BEGIN: for support llm-based hypothesis selection  -----
+    def _cosine_similarity_matrix_numpy(A, B):
+        dot_products = np.matmul(A, B.T)
+        A_norms = np.linalg.norm(A, axis=1, keepdims=True)
+        B_norms = np.linalg.norm(B, axis=1, keepdims=True).T
+        return dot_products / (A_norms * B_norms)
+
+    def _gumbel_softmax_hard_sample(logits, tau=1.0, n_samples=1):
+
+        gumbel_noise = -np.log(-np.log(np.random.uniform(size=logits.shape) + 1e-20) + 1e-20)
+        y = (logits + gumbel_noise) / tau
+        # softmax
+        y_soft = np.exp(y - np.max(y, axis=1, keepdims=True))
+        y_soft = y_soft / np.sum(y_soft, axis=1, keepdims=True)
+
+        sampled_indices = []
+        for i in range(y_soft.shape[0]):
+            choices = np.arange(y_soft.shape[1])
+            idx = np.random.choice(choices, size=n_samples, replace=False, p=y_soft[i])
+            sampled_indices.append(idx)
+        sampled_indices = np.unique(np.concatenate(sampled_indices))
+        return sampled_indices.tolist()
+
+    def _prob_dis(
+        self,
+        current_sota_score_in_current_trace,
+        extra_hypo_l: list[tuple[DSHypothesis, float]],
+        hypothesis_candidates,
+        competition,
+        path_length,
+    ):
+        # TODO: typing
+        history_hypo_str, history_scores = [], []
+        for hypo, score in extra_hypo_l:
+            history_hypo_str.append(hypo.hypothesis)
+            history_scores.append(score)
+
+        target_texts = [v["hypothesis"] for v in hypothesis_candidates.values()]
+        target_embs = np.array(APIBackend().create_embedding(target_texts), dtype=np.float32)
+
+        if not history_hypo_str:
+            return []
+        history_embs = np.array(APIBackend().create_embedding(history_hypo_str), dtype=np.float32)
+        # TODO: Here is an example to help understand the code:(Please check the correctness of the comment
+        # history_embs: numpy.ndarray of shape (N, D) where N is the number of historical hypotheses
+        # and D is the embedding dimension returned by APIBackend().create_embedding.
+        # It contains vector representations of each hypothesis string in history_hypo_str,
+        # used for computing similarity with target embeddings.
+        # Example: if history_hypo_str = ["Try RandomForest with 200 estimators", "Use LightGBM with early stopping"]
+        # and embedding dimension D=3, history_embs might be:
+        # array([[ 0.123, -0.456,  0.789],
+        #        [ 0.234,  0.567, -0.890]], dtype=float32)
+        sim_matrix = self._cosine_similarity_matrix_numpy(target_embs, history_embs)
+        candidate_scores = np.full((len(target_texts), 1), current_sota_score_in_current_trace, dtype=np.float32)
+        history_scores = np.array(history_scores, dtype=np.float32).reshape(1, -1)
+        bigger_is_better = get_metric_direction(competition)
+        if bigger_is_better:
+            score_diff_matrix = history_scores - candidate_scores
+        else:
+            score_diff_matrix = candidate_scores - history_scores
+        alpha, beta = 1.0, 1.0
+        if current_sota_score_in_current_trace == -1:  # FIXME: less magic number;
+            alpha, beta = 1.0, 0
+        gamma = math.log(2) / 30
+        logits = alpha * sim_matrix * math.exp(-gamma * path_length) + beta * np.tanh(score_diff_matrix)
+        logits_max = np.max(logits, axis=1, keepdims=True)
+        exp_logits = np.exp(logits - logits_max)
+        probs = exp_logits / np.sum(exp_logits, axis=1, keepdims=True)
+        num_candidates = probs.shape[-1]
+        n_samples = min(2, num_candidates)
+        flat_indices = self._gumbel_softmax_hard_sample(np.log(probs + 1e-20), tau=0.01, n_samples=n_samples)
+        if bigger_is_better:
+            best_idx = history_scores[0].argmax().item()
+            best_entry = (history_hypo_str[best_idx], history_scores[0, best_idx])
+        else:
+            best_idx = history_scores[0].argmin().item()
+            best_entry = (history_hypo_str[best_idx], history_scores[0, best_idx])
+        if len(flat_indices) > 2:
+            flat_indices = flat_indices[:2]
+        sampled_history_list = [best_entry] + [
+            (history_hypo_str[i], history_scores[0, i]) for i in flat_indices if i != best_idx
+        ]
+        return sampled_history_list
+
+    def _get_path(self, node, parent_nodes):
+        # FIXME: we should remove it in the future.
+        path = [node]
+        parent = parent_nodes.get(node)
+        if parent is not None:
+            path.extend(self._get_path(parent, parent_nodes))
+        return path
+
+    def _get_current_exp_score_list(self, trace, competition):
+        parent_nodes = {}
+        for node in range(len(trace.hist)):
+            parents = trace.get_parents(node)
+            parent_nodes[node] = parents[-2] if len(parents) > 1 else None
+        # FIXME: add the convert logic to method in trace
+        if hasattr(trace, "idx2loop_id"):
+            parent_nodes = {
+                trace.idx2loop_id[n]: trace.idx2loop_id[r] if r is not None else r for n, r in parent_nodes.items()
+            }
+        if trace.current_selection:
+            current_parent_record_id = trace.current_selection[0]  # record id
+        else:
+            return -1, 0
+        # current_parent_loop_id = trace.idx2loop_id[current_parent_record_id]# loop id
+        loop_id2idx = {v: k for k, v in trace.idx2loop_id.items()}
+
+        loop_id_list = self._get_path(trace.idx2loop_id[current_parent_record_id], parent_nodes)
+
+        score_list = [
+            trace.hist[loop_id2idx[loop_id]][0].result.loc["ensemble"].iloc[0].round(3)
+            for loop_id in loop_id_list
+            if trace.hist[loop_id2idx[loop_id]][1].decision == True
+        ]
+        if score_list:
+            bigger_is_better = get_metric_direction(competition)
+            if bigger_is_better:
+                return max(score_list), len(loop_id_list)
+            else:
+                return min(score_list), len(loop_id_list)
+        else:
+            return -1, len(loop_id_list)
+
+    def _llm_select_extra_hypo(self, trace: DSTrace) -> list[tuple[str, float]]:
+        """
+        Retrieve a list of additional hypotheses along with their ensemble scores
+        from the given experiment trace, intended for input into an LLM-based selection mechanism.
+
+        Parameters:
+            trace (DSTrace):
+
+        Returns:
+            list[tuple[str, float]]:
+                A list of tuples, where each tuple consists of:
+                    - str: The hypothesis description from a selected experiment.
+                      Example: "Use XGBoost with tuned learning_rate".
+                    - float: The associated ensemble result score, rounded to 3 decimal places.
+                      Example: 0.845
+                Example:
+                    [
+                        ("Try RandomForest with 200 estimators", 0.812),
+                        ("Use LightGBM with early stopping", 0.834)
+                    ]
+        """
+        return [
+            (exp.hypothesis, exp.result.loc["ensemble"].iloc[0])
+            for exp, _ in trace.experiment_and_feedback_list_after_init(return_type="sota", search_type="all")
+        ]
+
+    @wait_retry(retry_n=5)
+    def hypothesis_select_with_llm(
+        self,
+        scenario_desc: str,
+        exp_feedback_list_desc: str,
+        sota_exp_desc: str,
+        hypothesis_candidates: dict,
+        trace: DSTrace,
+    ):
+        res_time = RD_Agent_TIMER_wrapper.timer.remain_time()
+        ratio_merge_or_ensemble = DS_RD_SETTING.ratio_merge_or_ensemble
+
+        total_time = RD_Agent_TIMER_wrapper.timer.all_duration
+        # FIXME: total_time could be None
+        use_time = round(total_time.total_seconds(), 2) - round(res_time.total_seconds(), 2)
+        use_ratio = 100 * use_time / round(total_time.total_seconds(), 2)
+        use_ratio = round(use_ratio, 2)
+
+        full_time = self.scen.real_full_timeout() / 3600
+        # FIXME: less magic number
+        time_list_success = [-3600] + [
+            tr[0].running_info.running_time
+            for tr in trace.retrieve_search_list(search_type="ancestors")
+            if getattr(tr[1], "decision", False)
+        ]
+        time_max = max(time_list_success) / 3600
+        # sota_flag = (hasattr(trace, "sota_exp_to_submit") and trace.sota_exp_to_submit is not None)----> V10 CODE VERSION
+        bvs = BestValidSelector()  # ----> V14 CODE VERSION
+        sota_exp = bvs.get_sota_exp_to_submit(trace)  # ----> V14 CODE VERSION
+        sota_flag = sota_exp is not None and sota_exp.result is not None  # ----> V14 CODE VERSION
+
+        if sota_flag:
+            current_sota_score = sota_exp.result.loc["ensemble"].iloc[0].round(3)  # ----> V14 CODE VERSION
+            # trace.sota_exp_to_submit.result.loc["ensemble"].iloc[0].round(3) ----> V10 CODE VERSION
+        else:
+            current_sota_score = -1
+
+        competition = trace.scen.competition
+        if sota_flag:
+            current_sota_score_in_current_trace, path_length = self._get_current_exp_score_list(trace, competition)
+        else:
+            current_sota_score_in_current_trace = -1
+            path_length = 0
+
+        # extra_exp_feedback_list_desc: str,
+        # exp_feedback_scores: list,
+        extra_hypo_l = self._llm_select_extra_hypo(trace)
+        if len(extra_hypo_l) > 0:
+            # TODO:
+            selected_extra_hypo_l = self._prob_dis(
+                current_sota_score_in_current_trace,
+                extra_hypo_l,
+                hypothesis_candidates,
+                competition,
+                path_length,
+            )
+        else:
+            selected_extra_hypo_l = None
+        hypothesis_candidates = str(json.dumps(hypothesis_candidates, indent=2))
+
+        sys_prompt = T(".prompts_v2:hypothesis_select.system").r(
+            hypothesis_candidates=hypothesis_candidates,
+            res_time=round(res_time.total_seconds() / 3600, 2),
+            full_time=full_time,
+            use_ratio=use_ratio,
+            time_max=round(time_max, 2),
+            merge_hours=DS_RD_SETTING.merge_hours,
+            # extra_exp_feedback_list_desc=extra_exp_feedback_list_str,
+            selected_extra_hypo_l=selected_extra_hypo_l,
+            hypothesis_output_format=T(".prompts_v2:output_format.hypothesis_select_format").r(),
+            sota_flag=sota_flag,
+            current_sota_score=current_sota_score,
+            ratio_merge_or_ensemble=ratio_merge_or_ensemble,
+            current_sota_score_in_current_trace=current_sota_score_in_current_trace,
+        )
+
+        user_prompt = T(".prompts_v2:hypothesis_select.user").r(
+            scenario_desc=scenario_desc,
+            exp_and_feedback_list_desc=exp_feedback_list_desc,
+            sota_exp_desc=sota_exp_desc,
+        )
+
+        response = APIBackend().build_messages_and_create_chat_completion(
+            user_prompt=user_prompt,
+            system_prompt=sys_prompt,
+            response_format=HypothesisSimple if self.supports_response_schema else {"type": "json_object"},
+            json_target_type=(
+                Dict[str, Dict[str, str | Dict[str, str | int]]] if not self.supports_response_schema else None
+            ),
+        )
+
+        response_dict = json.loads(response)
+        assert response_dict.get("component") in HypothesisComponent.__members__, f"Invalid component"
+        assert response_dict.get("hypothesis") is not None, f"Invalid hypothesis"
+        return response_dict
+
+    # END: for support llm-based hypothesis selection  -----
+
    def hypothesis_rank(
        self, hypothesis_dict: dict, problem_dict: dict, selected_idx: Optional[int] = None
    ) -> Tuple[str, DSHypothesis]:
@@ -954,6 +1221,7 @@ class DSProposalV2ExpGen(ExpGen):
            workflow_check=workflow_check,
            metric_name=self.scen.metric_name,
            sibling_tasks=sibling_tasks,
+            fix_seed_and_data_split=DS_RD_SETTING.fix_seed_and_data_split,
        )
        user_prompt = T(".prompts_v2:task_gen.user").r(
            scenario_desc=scenario_desc,
@@ -1082,7 +1350,7 @@ class DSProposalV2ExpGen(ExpGen):
            type="failed",
            pipeline=pipeline,
        )
-
+        #
        # NOTE: we currently don't support inject diverse problems for the parallel + multi-trace mode,
        if DS_RD_SETTING.enable_inject_diverse and len(trace.hist) > 0:
            if len(trace.current_selection) == 0:
@@ -1120,6 +1388,9 @@ class DSProposalV2ExpGen(ExpGen):
                competition_desc=self.scen.get_competition_full_desc(),
            )

+        # sub-trace begin flag
+        is_new_tree = trace.is_selection_new_tree()
+
        # Step 2: Propose hypothesis based on the identified problems (and sampled ideas)
        hypothesis_dict = self.hypothesis_gen(
            component_desc=component_desc,
@@ -1131,6 +1402,7 @@ class DSProposalV2ExpGen(ExpGen):
            enable_idea_pool=DS_RD_SETTING.enable_knowledge_base,
            inject_diverse=inject_diverse,
            exp_gen_plan=plan.get("exp_gen") if plan else None,
+            is_new_tree=is_new_tree,
            packages_prompt=packages_prompt,
            sibling_exp=sibling_exp,
        )
@@ -1151,7 +1423,7 @@ class DSProposalV2ExpGen(ExpGen):
                    hypothesis_dict.pop(name)

        # Step 2.1 & 2.2: Hypothesis Critique and Rewrite Stage (controlled by enable_hypo_critique_rewrite)
-        if DS_RD_SETTING.enable_hypo_critique_rewrite:
+        if DS_RD_SETTING.enable_hypo_critique_rewrite and len(trace.hist) > 0:
            logger.info(f"Hypothesis critique and rewrite enabled - processing {len(hypothesis_dict)} hypotheses")

            # Critic Stage - Evaluate and identify flaws in hypotheses
@@ -1170,7 +1442,7 @@ class DSProposalV2ExpGen(ExpGen):

                # Rewriter Stage - Generate improved hypotheses based on critiques
                logger.info(f"Starting rewriter stage - generating improved hypotheses based on critique feedback")
-                improved_hypotheses_dict = self.hypothesis_rewrite(
+                hypothesis_dict = self.hypothesis_rewrite(
                    hypothesis_dict=hypothesis_dict,
                    critiques_dict=critiques_dict,
                    scenario_desc=scenario_desc,
@@ -1183,16 +1455,30 @@ class DSProposalV2ExpGen(ExpGen):
            except Exception as e:
                logger.warning(f"Hypothesis critique and rewrite failed: {e}")
                logger.info(f"Using original hypotheses as fallback instead of improved versions")
-                improved_hypotheses_dict = hypothesis_dict.copy()  # Use original hypotheses as fallback
        else:
            logger.info(f"Hypothesis critique and rewrite disabled - using original {len(hypothesis_dict)} hypotheses")
-            improved_hypotheses_dict = hypothesis_dict.copy()  # Use original hypotheses directly

        # Step 3: Select the best hypothesis
-        pickled_problem_name, new_hypothesis = self.hypothesis_rank(
-            hypothesis_dict=improved_hypotheses_dict,
-            problem_dict=all_problems,
-        )
+        if DS_RD_SETTING.llm_select_hypothesis:
+            response_dict = self.hypothesis_select_with_llm(
+                scenario_desc=scenario_desc,
+                exp_feedback_list_desc=exp_feedback_list_desc,
+                # extra_exp_feedback_list_desc=extra_exp_feedback_list_desc,
+                # exp_feedback_scores=exp_feedback_scores,
+                sota_exp_desc=sota_exp_desc,
+                hypothesis_candidates=hypothesis_dict,
+                trace=trace,
+            )
+            new_hypothesis = DSHypothesis(
+                component=response_dict.get("component"), hypothesis=response_dict.get("hypothesis")
+            )
+            pickled_problem_name = None
+        else:
+            pickled_problem_name, new_hypothesis = self.hypothesis_rank(
+                hypothesis_dict=hypothesis_dict,
+                problem_dict=all_problems,
+            )
+
        # Step 3.5: Update knowledge base with the picked problem
        if DS_RD_SETTING.enable_knowledge_base:
            trace.knowledge_base.update_pickled_problem(all_problems, pickled_problem_name)
@@ -1203,9 +1489,7 @@ class DSProposalV2ExpGen(ExpGen):
            sota_exp_desc=sota_exp_desc,
            sota_exp=sota_exp,
            hypotheses=(
-                [new_hypothesis]
-                if len(trace.hist) > 0
-                else self.get_all_hypotheses(all_problems, improved_hypotheses_dict)
+                [new_hypothesis] if len(trace.hist) > 0 else self.get_all_hypotheses(all_problems, hypothesis_dict)
            ),
            pipeline=pipeline,
            failed_exp_feedback_list_desc=failed_exp_feedback_list_desc,
@@ -119,14 +119,23 @@ class DataScienceScen(Scenario):
        )
        self.metric_name = response_json_analysis.get("Metric Name", "custom_metric")
        self.metric_direction_guess = response_json_analysis.get("Metric Direction", True)
-        self.longer_time_limit_required = (
+        # Determine if longer timeout is needed for coder and runner separately
+        base_longer_timeout_needed = (
            False
            if not DS_RD_SETTING.allow_longer_timeout
-            else (
-                response_json_analysis.get("Longer time limit required", False)
-                if DS_RD_SETTING.longer_timeout_by_llm
-                else True
-            )
+            else response_json_analysis.get("Longer time limit required", False)
+        )
+
+        self.coder_longer_time_limit_required = (
+            base_longer_timeout_needed
+            if DS_RD_SETTING.coder_enable_llm_decide_longer_timeout
+            else DS_RD_SETTING.allow_longer_timeout
+        )
+
+        self.runner_longer_time_limit_required = (
+            base_longer_timeout_needed
+            if DS_RD_SETTING.runner_enable_llm_decide_longer_timeout
+            else DS_RD_SETTING.allow_longer_timeout
        )

        # True or False, whether the competition scenario requires a longer time limit to the code.
@@ -136,9 +145,9 @@ class DataScienceScen(Scenario):
            DS_RD_SETTING.debug_timeout
            * min(
                DS_RD_SETTING.coder_longer_timeout_multiplier_upper,
-                self.timeout_increase_count * DS_RD_SETTING.timeout_increase_stage + 1,
+                self.timeout_increase_count * DS_RD_SETTING.coder_timeout_increase_stage + 1,
            )
-            if self.longer_time_limit_required
+            if self.coder_longer_time_limit_required
            else DS_RD_SETTING.debug_timeout
        )

@@ -150,9 +159,9 @@ class DataScienceScen(Scenario):
            DS_RD_SETTING.full_timeout
            * min(
                DS_RD_SETTING.runner_longer_timeout_multiplier_upper,
-                self.timeout_increase_count * DS_RD_SETTING.timeout_increase_stage + 1,
+                self.timeout_increase_count * DS_RD_SETTING.runner_timeout_increase_stage + 1,
            )
-            if self.longer_time_limit_required
+            if self.runner_longer_time_limit_required
            else DS_RD_SETTING.full_timeout
        )

@@ -52,7 +52,7 @@ describe: # some template to describe some object
    {% 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 }}    
    {% if exp_and_feedback[0].running_info.running_time is not none %}
-    Experiment Running Time: {{ exp_and_feedback[0].running_info.running_time }} seconds
+    Experiment Running Time: {{ (exp_and_feedback[0].running_info.running_time ) | round(1) }} seconds
    {% endif %}
    {% if exp_and_feedback[0].result is none %}
    Experiment Score: Running buggy