Improved Gepa UCB with range clamping / target

opto/features/gepa/gepa_algorithms.py

@@ -224,6 +224,14 @@ class GEPAUCBSearch(UCBSearchAlgorithm):
       - Optional periodic Merge crossover (uniform per-parameter) with desirability checks
     """
 
+    def _rank(self, raw: float) -> float:
+        """
+        If a score_range is provided (lo, hi), clamp the scalar score into that band.
+        This keeps UCB-like behavior numerically stable without changing external APIs.
+        """
+        if getattr(self, "score_range", None) is None or raw is None: return raw
+        lo, hi = self.score_range;  return float(min(hi, max(lo, raw)))
+
     def __init__(self,
                  agent,
                  optimizer=None,
@@ -270,6 +278,7 @@ def train(self,
               pareto_subset_size: int = 24,
               num_search_iterations: int = 120,
               train_batch_size: int = 2,
+              score_range: Optional[Tuple[float, float]] = None,
               merge_every: int = 6,
               log_frequency: Optional[int] = None,
               save_frequency: Optional[int] = None,
@@ -282,6 +291,7 @@ def train(self,
         num_threads = num_threads or self.num_threads
         log_frequency = log_frequency or 5
         validate_ds = validation_dataset or train_dataset
+        self.score_range = score_range # Optional score clamping band for mean-based selections
 
         # Fix a Pareto subset (small, stable) to compute per-instance vectors
         assert len(validate_ds["inputs"]) > 0, "Empty dataset."
@@ -296,10 +306,13 @@ def train(self,
 
         # Seed with current params
         base_params = {p: copy.deepcopy(p.data) for p in self.optimizer.parameters}
-        v0, m0 = self._evaluate_on_pareto(base_params, guide, num_threads=num_threads)
-        buffer.append(Candidate(params=base_params, eval_vector=v0, mean=m0, id=self._next_id(), ancestors=set()))
+        v0, m0_raw = self._evaluate_on_pareto(base_params, guide, num_threads=num_threads)
+        m0 = self._rank(m0_raw)
+        buffer.append(Candidate(params=base_params, eval_vector=v0, mean=m0, id=self._next_id(),
+                                ancestors=set(), meta={"raw_mean": m0_raw}))
         print_color(f"[GEPA] Seed candidate mean={m0:.4f}", "cyan")
 
+
         metrics = {"best_means": [], "new_child_means": [], "merge_accepts": 0, "total_merges": 0}
 
         for it in range(1, num_search_iterations + 1):
@@ -322,14 +335,16 @@ def train(self,
                 continue
 
             # Evaluate child on Pareto subset
-            child_vec, child_mean = self._evaluate_on_pareto(update_dict, guide, num_threads=num_threads)
+            child_vec, child_mean_raw = self._evaluate_on_pareto(update_dict, guide, num_threads=num_threads)
+            child_mean = self._rank(child_mean_raw)
             child = Candidate(params=update_dict,
                               eval_vector=child_vec,
                               mean=child_mean,
                               id=self._next_id(),
                               parent_ids=(parent.id,),
                               ancestors=set(parent.ancestors) | {parent.id},
-                              created_iter=it)
+                              created_iter=it,
+                              meta={"raw_mean": child_mean_raw})
             buffer.append(child)
             metrics["new_child_means"].append(child_mean)
             print_color(f"[GEPA] iter {it}: child mean={child_mean:.4f} (train-batch≈{train_batch_mean})", "green")
@@ -347,6 +362,10 @@ def train(self,
                 if merged is not None:
                     merged.id = self._next_id()
                     merged.created_iter = it
+                    # preserve raw and clamp to range for ranking/logging
+                    _raw = merged.mean
+                    merged.meta["raw_mean"] = _raw
+                    merged.mean = self._rank(_raw)
                     buffer.append(merged)
                     metrics["merge_accepts"] += 1
                     print_color(f"[GEPA] Merge accepted: mean={merged.mean:.4f}", "magenta")

tests/unit_tests/test_gepa_algorithms.py

@@ -0,0 +1,260 @@
+import math
+import os
+import random
+import re
+from typing import Any, Dict, List, Tuple
+
+import numpy as np
+import pytest
+
+# Provide a light stub for optional graphviz dependency to allow imports without system graphviz
+import sys, types
+if "graphviz" not in sys.modules:
+    sys.modules["graphviz"] = types.SimpleNamespace(Digraph=object)
+
+from opto.trace.modules import model as trace_model
+from opto.trace.nodes import node as trace_node
+from opto.optimizers.optoprime_v2 import OptoPrimeV2
+import pytest
+from opto.features.gepa.gepa_algorithms import (
+        GEPAAlgorithmBase,
+        GEPAUCBSearch,
+        GEPABeamPareto,
+        _compute_pareto_counts,
+        _pareto_sample,
+        _uniform_merge_params
+    )
+from opto.trainer.algorithms.beamsearch_algorithm import BeamsearchAlgorithm
+from opto.trainer.evaluators import evaluate
+from opto.trainer.guide import Guide
+from opto.utils.llm import DummyLLM
+
+
+class ExactMatchGuide(Guide):
+    """Simple guide: score=1 if response == reference, else 0."""
+
+    def get_feedback(self, query: Any, response: Any, reference: Any, **kwargs):
+        score = float(response == reference)
+        feedback = f"Score: {score}. Response: {response}. Reference: {reference}."
+        return score, feedback
+
+
+@trace_model
+class AddAgent:
+    """Toy agent: returns x + param."""
+
+    def __init__(self, param: int = 0):
+        self.param = trace_node(int(param), trainable=True)
+
+    def forward(self, x: int) -> int:
+        return x + self.param
+
+
+def make_dummy_llm(suggest_value: int) -> DummyLLM:
+    """Dummy LLM that parses the variable name from the prompt and suggests a fixed value.
+
+    Matches the default XML-like output format expected by OptoPrimeV2.
+    """
+
+    def _llm_callable(messages, **kwargs):
+        # Extract the variable name from the #Variables section in the prompt
+        problem = messages[1]["content"] if isinstance(messages, (list, tuple)) and len(messages) > 1 else ""
+        name_match = re.findall(r"<variable name=\"\s*(.*?)\" type=.*>", problem)
+        var_name = name_match[0] if name_match else "param"
+        return (
+            f"""
+            <reasoning> Dummy reasoning based on the input messages. </reasoning>
+            <variable>
+            <name> {var_name} </name>
+            <value> {suggest_value} </value>
+            </variable>
+            """
+        )
+
+    return DummyLLM(_llm_callable)
+
+
+def make_dataset(target_add: int, n: int = 8) -> Dict[str, List[int]]:
+    xs = list(range(n))
+    infos = [x + target_add for x in xs]
+    return {"inputs": xs, "infos": infos}
+
+
+def build_optimizer(agent: AddAgent, suggest_value: int) -> OptoPrimeV2:
+    return OptoPrimeV2(agent.parameters(), llm=make_dummy_llm(suggest_value))
+
+
+def test_pareto_counting_and_sampling():
+    # Construct mock candidates with per-instance eval vectors where each wins on one dimension
+    from types import SimpleNamespace
+
+    class Cand(SimpleNamespace):
+        pass
+
+    A = Cand(eval_vector=[1.0, 0.1], wins=0, mean=0.55)
+    B = Cand(eval_vector=[0.2, 1.1], wins=0, mean=0.65)
+    cands = [A, B]
+
+    _compute_pareto_counts(cands)
+    assert A.wins == 1 and B.wins == 1
+
+    rng = random.Random(0)
+    # With equal wins, both should be sampled with similar probability
+    picks = [
+        _pareto_sample([A, B], temperature=1.0, rng=rng) for _ in range(100)
+    ]
+    a_count = sum(p is A for p in picks)
+    b_count = sum(p is B for p in picks)
+    assert abs(a_count - b_count) < 40  # rough balance
+
+
+def test_uniform_merge_params_uses_both_parents():
+    # Use two ParameterNodes to exercise merging across keys
+    @trace_model
+    class TwoParam:
+        def __init__(self):
+            self.a = trace_node(1, trainable=True)
+            self.b = trace_node(2, trainable=True)
+
+        def forward(self, x):
+            return self.a + self.b + x
+
+    m = TwoParam()
+    a_params = {p: (10 if p.py_name.endswith("a") else 20) for p in m.parameters()}
+    b_params = {p: (100 if p.py_name.endswith("a") else 200) for p in m.parameters()}
+
+    rng = random.Random(123)
+    merged = _uniform_merge_params(a_params, b_params, rng)
+    # For each key, merged value should be chosen from either a_params or b_params
+    for k, v in merged.items():
+        assert v in (a_params[k], b_params[k])
+
+
+@pytest.mark.parametrize(
+    "algo_cls,train_kwargs",
+    [
+        (GEPAAlgorithmBase, {"num_iters": 8, "train_batch_size": 2, "merge_every": 2}),
+        (GEPAUCBSearch, {"num_search_iterations": 8, "train_batch_size": 2, "merge_every": 2}),
+        (GEPABeamPareto, {"num_search_iterations": 8, "train_batch_size": 2, "merge_every": 2}),
+    ],
+)
+def test_gepa_variants_converge_on_dummyllm(algo_cls, train_kwargs):
+    target_add = 5
+    ds = make_dataset(target_add, n=6)
+    agent = AddAgent(param=0)
+    optimizer = build_optimizer(agent, suggest_value=target_add)
+
+    algo = algo_cls(agent=agent, optimizer=optimizer, logger=None, num_threads=1)
+
+    # Prepare kwargs and include 'verbose' only if supported
+    import inspect
+    call_kwargs = dict(guide=ExactMatchGuide(), train_dataset=ds, pareto_subset_size=4, num_threads=1)
+    sig = inspect.signature(algo.train)
+    if 'validation_dataset' in sig.parameters:
+        call_kwargs['validation_dataset'] = ds
+    else:
+        call_kwargs['validate_dataset'] = ds
+    call_kwargs.update(train_kwargs)
+    if 'verbose' in sig.parameters:
+        call_kwargs['verbose'] = False
+
+    metrics, best = algo.train(**call_kwargs)
+
+    # Best mean on pareto subset should be perfect
+    assert isinstance(best, float)
+    assert best == pytest.approx(1.0, rel=0, abs=1e-6)
+    # Agent parameter should be updated to target_add
+    assert agent.param.data == target_add
+
+
+def test_numpy_seeding_reproducibility():
+    """Test that numpy seeding ensures reproducible behavior."""
+    target_add = 3
+    ds = make_dataset(target_add, n=4)
+
+    # Test with same seed
+    results = []
+    for seed in [123, 123]:  # Same seed twice
+        agent = AddAgent(param=0)
+        optimizer = build_optimizer(agent, suggest_value=target_add)
+        algo = GEPAAlgorithmBase(agent=agent, optimizer=optimizer, logger=None, num_threads=1, rng_seed=seed)
+
+        metrics, best = algo.train(
+            guide=ExactMatchGuide(),
+            train_dataset=ds,
+            validate_dataset=ds,
+            pareto_subset_size=3,
+            num_iters=2,
+            train_batch_size=1,
+            merge_every=2,
+            num_threads=1,
+        )
+        results.append((metrics, best, agent.param.data))
+
+    # Results should be identical with same seed
+    assert results[0][1] == results[1][1]  # Same best score
+    assert results[0][2] == results[1][2]  # Same final parameter
+
+    # Test with different seed
+    agent_diff = AddAgent(param=0)
+    optimizer_diff = build_optimizer(agent_diff, suggest_value=target_add)
+    algo_diff = GEPAAlgorithmBase(agent=agent_diff, optimizer=optimizer_diff, logger=None, num_threads=1, rng_seed=456)
+
+    metrics_diff, best_diff = algo_diff.train(
+        guide=ExactMatchGuide(),
+        train_dataset=ds,
+        validate_dataset=ds,
+        pareto_subset_size=3,
+        num_iters=2,
+        train_batch_size=1,
+        merge_every=2,
+        num_threads=1,
+    )
+
+    # Both should converge but the process might differ
+    # (though with DummyLLM behavior is very predictable)
+    assert best_diff == pytest.approx(1.0, rel=0, abs=1e-6)
+
+
+def test_thread_safety_with_sequential_fallback():
+    """Test that algorithms work correctly with sequential fallback when batch_run unavailable."""
+    target_add = 1
+    ds = make_dataset(target_add, n=2)
+    agent = AddAgent(param=0)
+    optimizer = build_optimizer(agent, suggest_value=target_add)
+
+    # Test with num_threads=1 (should use sequential)
+    algo = GEPAAlgorithmBase(agent=agent, optimizer=optimizer, logger=None, num_threads=1)
+    metrics, best = algo.train(
+        guide=ExactMatchGuide(),
+        train_dataset=ds,
+        validate_dataset=ds,
+        pareto_subset_size=2,
+        num_iters=2,
+        train_batch_size=1,
+        merge_every=2,
+        num_threads=1,
+    )
+
+    assert best == pytest.approx(1.0, rel=0, abs=1e-6)
+    assert agent.param.data == target_add
+
+    # Test with num_threads=2 (may use parallel or fallback to sequential)
+    agent2 = AddAgent(param=0)
+    optimizer2 = build_optimizer(agent2, suggest_value=target_add)
+    algo2 = GEPAAlgorithmBase(agent=agent2, optimizer=optimizer2, logger=None, num_threads=2)
+
+    metrics2, best2 = algo2.train(
+        guide=ExactMatchGuide(),
+        train_dataset=ds,
+        validate_dataset=ds,
+        pareto_subset_size=2,
+        num_iters=2,
+        train_batch_size=1,
+        merge_every=2,
+        num_threads=2,
+    )
+
+    assert best2 == pytest.approx(1.0, rel=0, abs=1e-6)
+    assert agent2.param.data == target_add
+