POLCA + Trace trainer

benchmarks/hf_qa/hf_qa_loader.py

@@ -218,6 +218,7 @@ def build_trace_problem(
     num_test: Optional[int] = None,
     initial_instructions: str = _DEFAULT_INSTRUCTIONS,
     framework: str = "trace",
+    num_eval_repeats: int = 1,
     **_kwargs: Any,
 ) -> Dict[str, Any]:
     """Build a Trace-Bench task bundle for an HF QA task.
@@ -271,14 +272,30 @@ def build_trace_problem(
         train_split = cfg.get("train_split", fallback)
         test_split  = cfg.get("test_split",  train_split)
 
-        train_val_tasks = _load_hf_data(cfg, train_split, num_train + num_validate)
-        train_tasks = train_val_tasks[:num_train]
-        val_tasks   = train_val_tasks[num_train:]
-        test_tasks  = _load_hf_data(cfg, test_split, num_test) if num_test > 0 else []
+        if train_split == test_split:
+            # Single-split mode (matches the reference hotpotqa_eval.py): load
+            # max(num_train, num_validate, num_test) examples once and let all
+            # three datasets take the *first N* — overlapping by design so val
+            # and test scores are directly comparable on identical samples.
+            n_load = max(num_train, num_validate, num_test)
+            all_tasks = _load_hf_data(cfg, train_split, n_load)
+            train_tasks = all_tasks[:num_train]
+            val_tasks   = all_tasks[:num_validate]
+            test_tasks  = all_tasks[:num_test] if num_test > 0 else []
+        else:
+            # Disjoint train/val + separate test split (legacy behavior).
+            train_val_tasks = _load_hf_data(cfg, train_split, num_train + num_validate)
+            train_tasks = train_val_tasks[:num_train]
+            val_tasks   = train_val_tasks[num_train:]
+            test_tasks  = _load_hf_data(cfg, test_split, num_test) if num_test > 0 else []
 
     objective = cfg.get("objective", "Optimize the agent's instructions for this QA task.")
     agent_class = cfg.get("agent_class", "hfqa")
 
+    # Allow hf_tasks.yaml to override the default initial instructions per-task.
+    if initial_instructions == _DEFAULT_INSTRUCTIONS:
+        initial_instructions = cfg.get("initial_instructions", _DEFAULT_INSTRUCTIONS)
+
     agent_kwargs: Dict[str, Any] = dict(
         initial_instructions=initial_instructions,
         node_description=cfg.get(
@@ -311,4 +328,6 @@ def build_trace_problem(
         bundle["validate_dataset"] = _make_dataset(val_tasks, agent_class)
     if test_tasks:
         bundle["test_dataset"] = _make_dataset(test_tasks, agent_class)
+    # Number of times to evaluate each example during dataset scoring (final eval only).
+    bundle["num_eval_repeats"] = max(1, int(num_eval_repeats))
     return bundle

benchmarks/hf_qa/hf_tasks.yaml

@@ -26,10 +26,10 @@
 #   objective       : optimizer objective injected into optimizer_kwargs
 
 - id: hotpot_qa
-  dataset_id: hotpot_qa
+  dataset_id: hotpotqa/hotpot_qa
   dataset_config: distractor
-  train_split: train        # train + validation examples come from the HF train split
-  test_split: validation    # test examples come from the HF validation split
+  train_split: validation   # train/val/test all from HF validation split (matches local hotpotqa_eval.py)
+  test_split: validation
   num_train: 100
   num_validate: 100
   num_test: 100

benchmarks/hf_qa/hotpot_qa.py

@@ -68,7 +68,23 @@ class HotpotQAGuide(Guide):
         """
 
         def get_feedback(self, task: Any, response: Any, info: Any, **kwargs: Any):
-            response_str = str(getattr(response, "data", response))
+            # Extract clean text. Trace LLM backends return AssistantTurn whose
+            # str() falls through to repr() ("AssistantTurn(content=..., model=...)")
+            # which pollutes substring matching. Prefer get_text() / .content if
+            # available, then .data, then plain str().
+            payload = getattr(response, "data", response)
+            if hasattr(payload, "get_text"):
+                response_str = str(payload.get_text())
+            elif hasattr(payload, "content"):
+                # AssistantTurn-like: content is ContentBlockList
+                response_str = str(payload.content)
+            else:
+                response_str = str(payload)
+            if not response_str.strip():
+                return None, (
+                    f"skipped: LLM returned empty response (likely safety filter or recitation).\n"
+                    f"Question: {info.question}"
+                )
             if check_answer(response_str, info.answer):
                 return 1.0, "Correct."
             return 0.0, (

configs/smallexp_hf_hotpotqa_gemini.yaml

@@ -0,0 +1,41 @@
+runs_dir: runs
+mode: real
+seeds: [42, 123, 456]
+
+job_timeout: 108000
+
+llm:
+  model: gemini-3.1-flash-lite-preview
+
+tasks:
+  - id: hf:hotpot_qa
+    eval_kwargs:
+      num_train: 100
+      num_validate: 100
+      num_test: 100
+      num_eval_repeats: 5
+
+trainers:
+  - id: POLCA
+    params_variants:
+      - ps_steps: 60
+        ps_batches: 1
+        ps_candidates: 5
+        ps_proposals: 1
+        batch_size: 2
+        num_threads: 10
+        test_frequency: null
+        log_frequency: 1
+        validate_exploration_candidates: false
+        use_best_candidate_to_explore: false
+        num_pool: 10              # post-training: re-evaluate top 10 candidates on full val, apply best
+        # POLCA-specific knobs
+        epsilon: 0.1
+        use_summarizer: true
+        embedding_model: gemini/gemini-embedding-001
+
+  - id: DSPyTrainer
+    params_variants:
+      - dspy_optimizer: gepa
+        auto: medium
+        reflection_minibatch_size: 10

configs/smoke_hf_hotpotqa_gemini.yaml

@@ -7,7 +7,7 @@ llm:
 
 tasks:
   - id: hf:hotpot_qa
-    eval_kwargs: {num_train: 20, num_validate: 10, num_test: 10}
+    eval_kwargs: {num_train: 1, num_validate: 1, num_test: 1}
 
 trainers:
   - id: PrioritySearch

trace_bench/cli.py

@@ -74,10 +74,30 @@ def _task_in_bench(task_key: str, bench: str | None) -> bool:
     "ps_candidates",
     "ps_proposals",
     "ps_mem_update",
+    "batch_size",
+    # POLCA-specific knobs (popped inside POLCA.train)
+    "epsilon",
+    "use_summarizer",
+    "embedding_model",
+    "num_pool",
+    # PrioritySearch fine-grained knobs (defaults are in PrioritySearch.train signature)
+    "validate_exploration_candidates",
+    "use_best_candidate_to_explore",
+    "test_frequency",
+    "num_test_samples",
+    "log_frequency",
+    "save_frequency",
     "gepa_iters",
     "gepa_train_bs",
     "gepa_merge_every",
     "gepa_pareto_subset",
+    # DSPyTrainer pass-through knobs
+    "dspy_optimizer",
+    "auto",
+    "max_steps",
+    "bsize",
+    "max_demos",
+    "reflection_minibatch_size",
     # LLM4AD pass-through knobs (merged into params_variants by config parser)
     "optimizer_kwargs",
     "eval_kwargs",

trace_bench/resolve.py

@@ -7,7 +7,7 @@
 
 
 def _default_trainer_kwargs(algo_name: str) -> Dict[str, Any]:
-    if algo_name == "PrioritySearch":
+    if algo_name in ("PrioritySearch", "POLCA"):
         return dict(num_epochs=1, num_steps=1, num_batches=1, num_candidates=2, num_proposals=2)
     if algo_name == "GEPA-Base":
         return dict(num_iters=1, train_batch_size=2, merge_every=2, pareto_subset_size=2)

trace_bench/runner.py

@@ -384,25 +384,42 @@ def _score_dataset(bundle: Dict[str, Any], dataset: Dict[str, Any], dataset_name
     if n <= 0:
         return {"score": None, "feedback": f"{dataset_name}: empty_dataset"}
 
-    scores: List[float] = []
+    repeats = max(1, int(bundle.get("num_eval_repeats", 1) or 1))
+
+    scores: List[float] = []  # one per (example, repeat) pair
     feedbacks: List[str] = []
     for i in range(n):
         task_input = inputs[i]
         task_info = infos[i]
-        response = _extract_response(bundle["param"], task_input)
-        try:
-            score, feedback = guide(task_input, response, task_info)
-        except Exception as exc:
-            return {"score": None, "feedback": f"{dataset_name}: eval_error[{i}]: {exc}"}
-        try:
-            scores.append(float(score))
-        except Exception:
-            return {"score": score, "feedback": f"{dataset_name}: non_numeric_score[{i}]: {feedback}"}
-        feedbacks.append(str(feedback))
+        per_example_scores: List[float] = []
+        last_feedback = ""
+        for r in range(repeats):
+            response = _extract_response(bundle["param"], task_input)
+            try:
+                score, feedback = guide(task_input, response, task_info)
+            except Exception as exc:
+                return {"score": None, "feedback": f"{dataset_name}: eval_error[{i},rep{r}]: {exc}"}
+            # Score=None means "unscoreable" (e.g. SAFETY filter); skip this rep.
+            if score is None:
+                last_feedback = str(feedback)
+                continue
+            try:
+                s = float(score)
+            except Exception:
+                return {"score": score, "feedback": f"{dataset_name}: non_numeric_score[{i},rep{r}]: {feedback}"}
+            per_example_scores.append(s)
+            scores.append(s)
+            last_feedback = str(feedback)
+        feedbacks.append(last_feedback)
 
+    if not scores:
+        return {
+            "score": None,
+            "feedback": f"{dataset_name}: all {n} examples x {repeats} repeats produced unscoreable responses.",
+        }
     return {
         "score": sum(scores) / len(scores),
-        "feedback": f"{dataset_name}: mean over {len(scores)} examples. " + " | ".join(feedbacks[:3]),
+        "feedback": f"{dataset_name}: mean over {len(scores)} valid evals (out of {n} examples x {repeats} repeats). " + " | ".join(feedbacks[:3]),
     }