diff --git a/rdagent/scenarios/kaggle/experiment/forest-cover-type-prediction_template/train.py b/rdagent/scenarios/kaggle/experiment/forest-cover-type-prediction_template/train.py
index 104cb4c7..458b3301 100644
--- a/rdagent/scenarios/kaggle/experiment/forest-cover-type-prediction_template/train.py
+++ b/rdagent/scenarios/kaggle/experiment/forest-cover-type-prediction_template/train.py
@@ -7,7 +7,6 @@ import pandas as pd
 from fea_share_preprocess import clean_and_impute_data, preprocess_script
 from scipy import stats
 from sklearn.metrics import accuracy_score, matthews_corrcoef
-from sklearn.preprocessing import LabelEncoder
 
 # Set random seed for reproducibility
 SEED = 42
@@ -37,7 +36,6 @@ def import_module_from_path(module_name, module_path):
 
 
 # 1) Preprocess the data
-# TODO 如果已经做过数据预处理了，不需要再做了
 X_train, X_valid, y_train, y_valid, X_test, ids = preprocess_script()
 
 # 2) Auto feature engineering
diff --git a/rdagent/scenarios/kaggle/experiment/forest-cover-type-prediction_template/model/model_nn.py b/rdagent/scenarios/kaggle/experiment/meta_tpl/model/model_nn.py
similarity index 65%
rename from rdagent/scenarios/kaggle/experiment/forest-cover-type-prediction_template/model/model_nn.py
rename to rdagent/scenarios/kaggle/experiment/meta_tpl/model/model_nn.py
index 53214626..91b912ac 100644
--- a/rdagent/scenarios/kaggle/experiment/forest-cover-type-prediction_template/model/model_nn.py
+++ b/rdagent/scenarios/kaggle/experiment/meta_tpl/model/model_nn.py
@@ -9,53 +9,52 @@ from tqdm import tqdm
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 
-# Modified model for multi-class classification
-class HybridFeatureInteractionModel(nn.Module):
-    def __init__(self, num_features, num_classes):
-        super(HybridFeatureInteractionModel, self).__init__()
+# Restored three-layer model structure
+class FeatureInteractionModel(nn.Module):
+    def __init__(self, num_features):
+        super(FeatureInteractionModel, self).__init__()
         self.fc1 = nn.Linear(num_features, 128)
         self.bn1 = nn.BatchNorm1d(128)
         self.fc2 = nn.Linear(128, 64)
         self.bn2 = nn.BatchNorm1d(64)
-        self.fc3 = nn.Linear(64, num_classes)  # Output nodes equal to num_classes
+        self.fc3 = nn.Linear(64, 1)
         self.dropout = nn.Dropout(0.3)
 
     def forward(self, x):
         x = F.relu(self.bn1(self.fc1(x)))
         x = F.relu(self.bn2(self.fc2(x)))
         x = self.dropout(x)
-        x = self.fc3(x)  # No activation here, use CrossEntropyLoss
+        x = torch.sigmoid(self.fc3(x))
         return x
 
 
 # Training function
 def fit(X_train, y_train, X_valid, y_valid):
     num_features = X_train.shape[1]
-    num_classes = len(np.unique(y_train))  # Determine number of classes
-    model = HybridFeatureInteractionModel(num_features, num_classes).to(device)
-    criterion = nn.CrossEntropyLoss()  # Use CrossEntropyLoss for multi-class
+    model = FeatureInteractionModel(num_features).to(device)
+    criterion = nn.BCELoss()  # Binary classification problem
     optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
 
     # Convert to TensorDataset and create DataLoader
     train_dataset = TensorDataset(
-        torch.tensor(X_train.to_numpy(), dtype=torch.float32), torch.tensor(y_train.to_numpy(), dtype=torch.long)
+        torch.tensor(X_train.to_numpy(), dtype=torch.float32), torch.tensor(y_train.reshape(-1), dtype=torch.float32)
     )
     valid_dataset = TensorDataset(
-        torch.tensor(X_valid.to_numpy(), dtype=torch.float32), torch.tensor(y_valid.to_numpy(), dtype=torch.long)
+        torch.tensor(X_valid.to_numpy(), dtype=torch.float32), torch.tensor(y_valid.reshape(-1), dtype=torch.float32)
     )
     train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
     valid_loader = DataLoader(valid_dataset, batch_size=32, shuffle=False)
 
     # Train the model
     model.train()
-    for epoch in range(5):  # just for quick run
+    for epoch in range(5):
         print(f"Epoch {epoch + 1}/5")
         epoch_loss = 0
         for X_batch, y_batch in tqdm(train_loader, desc="Training", leave=False):
-            X_batch, y_batch = X_batch.to(device), y_batch.to(device)
+            X_batch, y_batch = X_batch.to(device), y_batch.to(device)  # Move data to the device
             optimizer.zero_grad()
-            outputs = model(X_batch)
-            loss = criterion(outputs, y_batch)
+            outputs = model(X_batch).squeeze(1)  # Reshape outputs to [32]
+            loss = criterion(outputs, y_batch)  # Adjust target shape
             loss.backward()
             optimizer.step()
             epoch_loss += loss.item()
@@ -69,10 +68,9 @@ def predict(model, X):
     model.eval()
     predictions = []
     with torch.no_grad():
-        X_tensor = torch.tensor(X.values, dtype=torch.float32).to(device)
+        X_tensor = torch.tensor(X.values, dtype=torch.float32).to(device)  # Move data to the device
         for i in tqdm(range(0, len(X_tensor), 32), desc="Predicting", leave=False):
-            batch = X_tensor[i : i + 32]
-            pred = model(batch)
-            pred = torch.argmax(pred, dim=1).cpu().numpy()  # Use argmax to get class
+            batch = X_tensor[i : i + 32]  # Predict in batches
+            pred = model(batch).squeeze().cpu().numpy()  # Move results back to CPU
             predictions.extend(pred)
-    return np.array(predictions)
+    return np.array(predictions)  # Return boolean predictions
diff --git a/rdagent/scenarios/kaggle/experiment/playground-series-s4e8_template/model/model_nn.py b/rdagent/scenarios/kaggle/experiment/playground-series-s4e8_template/model/model_nn.py
index 00431100..91b912ac 100644
--- a/rdagent/scenarios/kaggle/experiment/playground-series-s4e8_template/model/model_nn.py
+++ b/rdagent/scenarios/kaggle/experiment/playground-series-s4e8_template/model/model_nn.py
@@ -10,9 +10,9 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 
 # Restored three-layer model structure
-class HybridFeatureInteractionModel(nn.Module):
+class FeatureInteractionModel(nn.Module):
     def __init__(self, num_features):
-        super(HybridFeatureInteractionModel, self).__init__()
+        super(FeatureInteractionModel, self).__init__()
         self.fc1 = nn.Linear(num_features, 128)
         self.bn1 = nn.BatchNorm1d(128)
         self.fc2 = nn.Linear(128, 64)
@@ -31,7 +31,7 @@ class HybridFeatureInteractionModel(nn.Module):
 # Training function
 def fit(X_train, y_train, X_valid, y_valid):
     num_features = X_train.shape[1]
-    model = HybridFeatureInteractionModel(num_features).to(device)
+    model = FeatureInteractionModel(num_features).to(device)
     criterion = nn.BCELoss()  # Binary classification problem
     optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
 
diff --git a/rdagent/scenarios/kaggle/experiment/playground-series-s4e9_template/model/model_nn.py b/rdagent/scenarios/kaggle/experiment/playground-series-s4e9_template/model/model_nn.py
index b45175e8..0fd122df 100644
--- a/rdagent/scenarios/kaggle/experiment/playground-series-s4e9_template/model/model_nn.py
+++ b/rdagent/scenarios/kaggle/experiment/playground-series-s4e9_template/model/model_nn.py
@@ -10,9 +10,9 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 
 # Modified model for regression
-class HybridFeatureInteractionModel(nn.Module):
+class FeatureInteractionModel(nn.Module):
     def __init__(self, num_features):
-        super(HybridFeatureInteractionModel, self).__init__()
+        super(FeatureInteractionModel, self).__init__()
         self.fc1 = nn.Linear(num_features, 128)
         self.bn1 = nn.BatchNorm1d(128)
         self.fc2 = nn.Linear(128, 64)
@@ -31,7 +31,7 @@ class HybridFeatureInteractionModel(nn.Module):
 # Training function
 def fit(X_train, y_train, X_valid, y_valid):
     num_features = X_train.shape[1]
-    model = HybridFeatureInteractionModel(num_features).to(device)
+    model = FeatureInteractionModel(num_features).to(device)
     criterion = nn.MSELoss()  # Use MSELoss for regression
     optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
 
diff --git a/rdagent/scenarios/kaggle/experiment/scenario.py b/rdagent/scenarios/kaggle/experiment/scenario.py
index 95fd3158..5cfcfcde 100644
--- a/rdagent/scenarios/kaggle/experiment/scenario.py
+++ b/rdagent/scenarios/kaggle/experiment/scenario.py
@@ -35,6 +35,7 @@ class KGScenario(Scenario):
         self.target_description = None
         self.competition_features = None
         self._analysis_competition_description()
+        self.if_action_choosing_based_on_UCB = False
 
         self._background = self.background
 
diff --git a/rdagent/scenarios/kaggle/experiment/sf-crime_template/model/model_nn.py b/rdagent/scenarios/kaggle/experiment/sf-crime_template/model/model_nn.py
new file mode 100644
index 00000000..91b912ac
--- /dev/null
+++ b/rdagent/scenarios/kaggle/experiment/sf-crime_template/model/model_nn.py
@@ -0,0 +1,76 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import DataLoader, TensorDataset
+from tqdm import tqdm
+
+# Check if a GPU is available
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+# Restored three-layer model structure
+class FeatureInteractionModel(nn.Module):
+    def __init__(self, num_features):
+        super(FeatureInteractionModel, self).__init__()
+        self.fc1 = nn.Linear(num_features, 128)
+        self.bn1 = nn.BatchNorm1d(128)
+        self.fc2 = nn.Linear(128, 64)
+        self.bn2 = nn.BatchNorm1d(64)
+        self.fc3 = nn.Linear(64, 1)
+        self.dropout = nn.Dropout(0.3)
+
+    def forward(self, x):
+        x = F.relu(self.bn1(self.fc1(x)))
+        x = F.relu(self.bn2(self.fc2(x)))
+        x = self.dropout(x)
+        x = torch.sigmoid(self.fc3(x))
+        return x
+
+
+# Training function
+def fit(X_train, y_train, X_valid, y_valid):
+    num_features = X_train.shape[1]
+    model = FeatureInteractionModel(num_features).to(device)
+    criterion = nn.BCELoss()  # Binary classification problem
+    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
+
+    # Convert to TensorDataset and create DataLoader
+    train_dataset = TensorDataset(
+        torch.tensor(X_train.to_numpy(), dtype=torch.float32), torch.tensor(y_train.reshape(-1), dtype=torch.float32)
+    )
+    valid_dataset = TensorDataset(
+        torch.tensor(X_valid.to_numpy(), dtype=torch.float32), torch.tensor(y_valid.reshape(-1), dtype=torch.float32)
+    )
+    train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
+    valid_loader = DataLoader(valid_dataset, batch_size=32, shuffle=False)
+
+    # Train the model
+    model.train()
+    for epoch in range(5):
+        print(f"Epoch {epoch + 1}/5")
+        epoch_loss = 0
+        for X_batch, y_batch in tqdm(train_loader, desc="Training", leave=False):
+            X_batch, y_batch = X_batch.to(device), y_batch.to(device)  # Move data to the device
+            optimizer.zero_grad()
+            outputs = model(X_batch).squeeze(1)  # Reshape outputs to [32]
+            loss = criterion(outputs, y_batch)  # Adjust target shape
+            loss.backward()
+            optimizer.step()
+            epoch_loss += loss.item()
+        print(f"End of epoch {epoch + 1}, Avg Loss: {epoch_loss / len(train_loader):.4f}")
+
+    return model
+
+
+# Prediction function
+def predict(model, X):
+    model.eval()
+    predictions = []
+    with torch.no_grad():
+        X_tensor = torch.tensor(X.values, dtype=torch.float32).to(device)  # Move data to the device
+        for i in tqdm(range(0, len(X_tensor), 32), desc="Predicting", leave=False):
+            batch = X_tensor[i : i + 32]  # Predict in batches
+            pred = model(batch).squeeze().cpu().numpy()  # Move results back to CPU
+            predictions.extend(pred)
+    return np.array(predictions)  # Return boolean predictions
diff --git a/rdagent/scenarios/kaggle/experiment/spaceship-titanic_template/fea_share_preprocess.py b/rdagent/scenarios/kaggle/experiment/spaceship-titanic_template/fea_share_preprocess.py
index 02536382..8dcb2a3f 100644
--- a/rdagent/scenarios/kaggle/experiment/spaceship-titanic_template/fea_share_preprocess.py
+++ b/rdagent/scenarios/kaggle/experiment/spaceship-titanic_template/fea_share_preprocess.py
@@ -5,7 +5,7 @@ from sklearn.compose import ColumnTransformer
 from sklearn.impute import SimpleImputer
 from sklearn.model_selection import train_test_split
 from sklearn.pipeline import Pipeline
-from sklearn.preprocessing import LabelEncoder, OneHotEncoder
+from sklearn.preprocessing import LabelEncoder
 
 
 def prepreprocess():
@@ -17,7 +17,7 @@ def prepreprocess():
     data_df = data_df.drop(["PassengerId"], axis=1)
 
     X = data_df.drop(["Transported"], axis=1)
-    y = data_df[["Transported"]]
+    y = data_df["Transported"]
 
     label_encoder = LabelEncoder()
     y = label_encoder.fit_transform(y)  # Convert class labels to numeric
@@ -37,45 +37,39 @@ def preprocess_fit(X_train: pd.DataFrame):
     categorical_cols = [cname for cname in X_train.columns if X_train[cname].dtype == "object"]
 
     # Define preprocessors for numerical and categorical features
-    categorical_transformer = Pipeline(
-        steps=[
-            ("imputer", SimpleImputer(strategy="most_frequent")),
-            ("onehot", OneHotEncoder(handle_unknown="ignore")),
-        ]
-    )
+    label_encoders = {col: LabelEncoder().fit(X_train[col]) for col in categorical_cols}
 
     numerical_transformer = Pipeline(steps=[("imputer", SimpleImputer(strategy="mean"))])
 
     # Combine preprocessing steps
     preprocessor = ColumnTransformer(
         transformers=[
-            ("cat", categorical_transformer, categorical_cols),
             ("num", numerical_transformer, numerical_cols),
-        ]
+        ],
+        remainder="passthrough",
     )
 
     # Fit the preprocessor on the training data
     preprocessor.fit(X_train)
 
-    return preprocessor
+    return preprocessor, label_encoders
 
 
-def preprocess_transform(X: pd.DataFrame, preprocessor):
+def preprocess_transform(X: pd.DataFrame, preprocessor, label_encoders):
     """
     Transforms the given DataFrame using the fitted preprocessor.
     Ensures the processed data has consistent features across train, validation, and test sets.
     """
-    # Transform the data using the fitted preprocessor
-    X_array = preprocessor.transform(X).toarray()
+    # Encode categorical features
+    for col, le in label_encoders.items():
+        # Handle unseen labels by setting them to a default value (e.g., -1)
+        X[col] = X[col].apply(lambda x: le.transform([x])[0] if x in le.classes_ else -1)
 
-    # Get feature names for the columns in the transformed data
-    categorical_cols = [cname for cname in X.columns if X[cname].dtype == "object"]
-    feature_names = preprocessor.named_transformers_["cat"]["onehot"].get_feature_names_out(
-        categorical_cols
-    ).tolist() + [cname for cname in X.columns if X[cname].dtype in ["int64", "float64"]]
+    # Transform the data using the fitted preprocessor
+    X_array = preprocessor.transform(X)
 
     # Convert arrays back to DataFrames
-    X_transformed = pd.DataFrame(X_array, columns=feature_names, index=X.index)
+    X_transformed = pd.DataFrame(X_array, columns=X.columns, index=X.index)
 
     return X_transformed
 
@@ -96,16 +90,16 @@ def preprocess_script():
     X_train, X_valid, y_train, y_valid = prepreprocess()
 
     # Fit the preprocessor on the training data
-    preprocessor = preprocess_fit(X_train)
+    preprocessor, label_encoders = preprocess_fit(X_train)
 
     # Preprocess the train, validation, and test data
-    X_train = preprocess_transform(X_train, preprocessor)
-    X_valid = preprocess_transform(X_valid, preprocessor)
+    X_train = preprocess_transform(X_train, preprocessor, label_encoders)
+    X_valid = preprocess_transform(X_valid, preprocessor, label_encoders)
 
     # Load and preprocess the test data
     submission_df = pd.read_csv("/kaggle/input/test.csv")
     passenger_ids = submission_df["PassengerId"]
     submission_df = submission_df.drop(["PassengerId"], axis=1)
-    X_test = preprocess_transform(submission_df, preprocessor)
+    X_test = preprocess_transform(submission_df, preprocessor, label_encoders)
 
     return X_train, X_valid, y_train, y_valid, X_test, passenger_ids
diff --git a/rdagent/scenarios/kaggle/experiment/spaceship-titanic_template/model/model_nn.py b/rdagent/scenarios/kaggle/experiment/spaceship-titanic_template/model/model_nn.py
index 00431100..91b912ac 100644
--- a/rdagent/scenarios/kaggle/experiment/spaceship-titanic_template/model/model_nn.py
+++ b/rdagent/scenarios/kaggle/experiment/spaceship-titanic_template/model/model_nn.py
@@ -10,9 +10,9 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 
 # Restored three-layer model structure
-class HybridFeatureInteractionModel(nn.Module):
+class FeatureInteractionModel(nn.Module):
     def __init__(self, num_features):
-        super(HybridFeatureInteractionModel, self).__init__()
+        super(FeatureInteractionModel, self).__init__()
         self.fc1 = nn.Linear(num_features, 128)
         self.bn1 = nn.BatchNorm1d(128)
         self.fc2 = nn.Linear(128, 64)
@@ -31,7 +31,7 @@ class HybridFeatureInteractionModel(nn.Module):
 # Training function
 def fit(X_train, y_train, X_valid, y_valid):
     num_features = X_train.shape[1]
-    model = HybridFeatureInteractionModel(num_features).to(device)
+    model = FeatureInteractionModel(num_features).to(device)
     criterion = nn.BCELoss()  # Binary classification problem
     optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
 
diff --git a/rdagent/scenarios/kaggle/experiment/spaceship-titanic_template/model/model_xgboost.py b/rdagent/scenarios/kaggle/experiment/spaceship-titanic_template/model/model_xgboost.py
index a70fa680..64784224 100644
--- a/rdagent/scenarios/kaggle/experiment/spaceship-titanic_template/model/model_xgboost.py
+++ b/rdagent/scenarios/kaggle/experiment/spaceship-titanic_template/model/model_xgboost.py
@@ -22,7 +22,7 @@ def fit(X_train: pd.DataFrame, y_train: pd.DataFrame, X_valid: pd.DataFrame, y_v
     params = {
         "nthred": -1,
     }
-    num_round = 180
+    num_round = 100
 
     evallist = [(dtrain, "train"), (dvalid, "eval")]
     bst = xgb.train(params, dtrain, num_round, evallist)
diff --git a/rdagent/scenarios/kaggle/prompts.yaml b/rdagent/scenarios/kaggle/prompts.yaml
index c32cc25c..9494b984 100644
--- a/rdagent/scenarios/kaggle/prompts.yaml
+++ b/rdagent/scenarios/kaggle/prompts.yaml
@@ -39,7 +39,7 @@ hypothesis_and_feedback: |-
 hypothesis_output_format: |-
   The output should follow JSON format. The schema is as follows:
   {
-    "action": "The action that the user wants to take based on the information provided. should be one of ["Feature engineering", "Feature processing", "Model feature selection", "Model tuning"]"
+    "action": "If "hypothesis_specification" provides the action you need to take, please follow "hypothesis_specification" to choose the action. Otherwise, based on previous experimental results, suggest the action you believe is most appropriate at the moment. It should be one of ["Feature engineering", "Feature processing", "Model feature selection", "Model tuning"]"
     "hypothesis": "The new hypothesis generated based on the information provided.",
     "reason": "The reason why you generate this hypothesis. It should be comprehensive and logical. It should cover the other keys below and extend them.",
     "concise_reason": "Two-line summary. First line focuses on a concise justification for the change. Second line generalizes a knowledge statement.",
diff --git a/rdagent/scenarios/kaggle/proposal/proposal.py b/rdagent/scenarios/kaggle/proposal/proposal.py
index b0abe9be..0765c0a9 100644
--- a/rdagent/scenarios/kaggle/proposal/proposal.py
+++ b/rdagent/scenarios/kaggle/proposal/proposal.py
@@ -1,4 +1,5 @@
 import json
+import math
 from pathlib import Path
 from typing import List, Tuple
 
@@ -81,6 +82,20 @@ class KGHypothesisGen(ModelHypothesisGen):
 
     def __init__(self, scen: Scenario) -> Tuple[dict, bool]:
         super().__init__(scen)
+        self.action_counts = {
+            "Feature engineering": 0,
+            "Feature processing": 0,
+            "Model feature selection": 0,
+            "Model tuning": 0,
+        }
+        self.reward_estimates = {
+            "Feature engineering": 0.0,
+            "Feature processing": 0.0,
+            "Model feature selection": 0.0,
+            "Model tuning": 0.0,
+        }
+        self.confidence_parameter = 1.0
+        self.initial_performance = 0.0
 
     def generate_RAG_content(self, trace: Trace) -> str:
         if trace.knowledge_base is None:
@@ -162,6 +177,52 @@ class KGHypothesisGen(ModelHypothesisGen):
         )
         return RAG_content
 
+    def update_reward_estimates(self, trace: Trace) -> None:
+        if len(trace.hist) > 0:
+            last_entry = trace.hist[-1]
+            last_action = last_entry[0].action
+            last_result = last_entry[1].result
+            # Extract performance_t
+            performance_t = last_result.get("performance", 0.0)
+            # Get performance_{t-1}
+            if len(trace.hist) > 1:
+                prev_entry = trace.hist[-2]
+                prev_result = prev_entry[1].result
+                performance_t_minus_1 = prev_result.get("performance", 0.0)
+            else:
+                performance_t_minus_1 = self.initial_performance
+
+            reward = performance_t_minus_1 - performance_t
+            n_o = self.action_counts[last_action]
+            mu_o = self.reward_estimates[last_action]
+            self.reward_estimates[last_action] += (reward - mu_o) / n_o
+        else:
+            # First iteration, nothing to update
+            pass
+
+    def execute_next_action(self, trace: Trace) -> str:
+        actions = list(self.action_counts.keys())
+        t = sum(self.action_counts.values()) + 1
+
+        # If any action has not been tried yet, select it
+        for action in actions:
+            if self.action_counts[action] == 0:
+                selected_action = action
+                self.action_counts[selected_action] += 1
+                return selected_action
+
+        c = self.confidence_parameter
+        ucb_values = {}
+        for action in actions:
+            mu_o = self.reward_estimates[action]
+            n_o = self.action_counts[action]
+            ucb = mu_o + c * math.sqrt(math.log(t) / n_o)
+            ucb_values[action] = ucb
+        # Select action with highest UCB
+        selected_action = max(ucb_values, key=ucb_values.get)
+        self.action_counts[selected_action] += 1
+        return selected_action
+
     def prepare_context(self, trace: Trace) -> Tuple[dict, bool]:
         hypothesis_and_feedback = (
             (
@@ -173,16 +234,22 @@ class KGHypothesisGen(ModelHypothesisGen):
             else "No previous hypothesis and feedback available since it's the first round."
         )
 
+        if self.scen.if_action_choosing_based_on_UCB:
+            action = self.execute_next_action(trace)
+
         context_dict = {
             "hypothesis_and_feedback": hypothesis_and_feedback,
             "RAG": self.generate_RAG_content(trace),
             "hypothesis_output_format": prompt_dict["hypothesis_output_format"],
-            "hypothesis_specification": None,
+            "hypothesis_specification": f"next experiment action is {action}"
+            if self.scen.if_action_choosing_based_on_UCB
+            else None,
         }
         return context_dict, True
 
     def convert_response(self, response: str) -> ModelHypothesis:
         response_dict = json.loads(response)
+
         hypothesis = KGHypothesis(
             hypothesis=response_dict["hypothesis"],
             reason=response_dict["reason"],