Integrate GEPA via rl --gepa with acceptance gate and Pareto selection

configs/reverse_text/gepa.toml

@@ -0,0 +1,34 @@
+[model]
+name = "Qwen/Qwen3-0.6B"
+
+[evaluate]
+benchmarks = ["math500"]
+subset_size = 32
+rollouts_per_prompt = 1
+max_tokens = 64
+min_tokens = 0
+
+[operators]
+mutation_rate = 0.7
+crossover_rate = 0.3
+max_prompt_chars = 2000
+enforce_diversity = true
+min_levenshtein_distance = 40
+
+[selection]
+strategy = "top-k"
+k = 6
+keep_elite = 2
+
+population_size = 8
+generations = 3
+seed = 42
+
+[log]
+level = "info"
+
+[monitor]
+# Enable W&B by adding a [monitor.wandb] section externally if desired
+system_log_frequency = 0
+
+

environments/vf_acereason_math/vf_acereason_math.py

@@ -8,6 +8,7 @@ def load_environment(
     solve_rate_field: str | None = None,
     min_solve_rate: float | None = None,
     max_solve_rate: float | None = None,
+    system_prompt: str | None = None,
     **kwargs,
 ) -> vf.Environment:
     train_dataset = load_dataset("nvidia/AceReason-Math", split="train").map(
@@ -30,5 +31,9 @@ def correct_answer_reward_func(completion, info, **kwargs) -> float:
         weights=[1.0],
     )
 
-    vf_env = vf.SingleTurnEnv(dataset=train_dataset, rubric=rubric)
+    vf_env = (
+        vf.SingleTurnEnv(dataset=train_dataset, rubric=rubric, system_prompt=system_prompt)
+        if system_prompt is not None
+        else vf.SingleTurnEnv(dataset=train_dataset, rubric=rubric)
+    )
     return vf_env

environments/vf_deepscaler_math/vf_deepscaler_math.py

@@ -8,6 +8,7 @@ def load_environment(
     solve_rate_field: str | None = None,
     min_solve_rate: float | None = None,
     max_solve_rate: float | None = None,
+    system_prompt: str | None = None,
     **kwargs,
 ) -> vf.Environment:
     train_dataset = load_dataset("agentica-org/DeepScaleR-Preview-Dataset", split="train").map(
@@ -30,5 +31,9 @@ def correct_answer_reward_func(completion, info, **kwargs) -> float:
         weights=[1.0],
     )
 
-    vf_env = vf.SingleTurnEnv(dataset=train_dataset, rubric=rubric)
+    vf_env = (
+        vf.SingleTurnEnv(dataset=train_dataset, rubric=rubric, system_prompt=system_prompt)
+        if system_prompt is not None
+        else vf.SingleTurnEnv(dataset=train_dataset, rubric=rubric)
+    )
     return vf_env

environments/vf_hendrycks_math/vf_hendrycks_math.py

@@ -2,7 +2,7 @@
 from datasets import load_dataset
 
 
-def load_environment(**kwargs) -> vf.Environment:
+def load_environment(system_prompt: str | None = None, **kwargs) -> vf.Environment:
     import json
 
     from verifiers.utils.data_utils import extract_boxed_answer
@@ -31,5 +31,9 @@ def correct_answer_reward_func(completion, info) -> float:
         weights=[1.0],
     )
 
-    vf_env = vf.SingleTurnEnv(dataset=train_dataset, parser=parser, rubric=rubric)
+    # Pass optional system prompt if provided (SingleTurnEnv supports it)
+    if system_prompt is not None:
+        vf_env = vf.SingleTurnEnv(dataset=train_dataset, parser=parser, rubric=rubric, system_prompt=system_prompt)
+    else:
+        vf_env = vf.SingleTurnEnv(dataset=train_dataset, parser=parser, rubric=rubric)
     return vf_env

environments/vf_intellect_math/vf_intellect_math.py

@@ -6,6 +6,7 @@ def load_environment(
     solve_rate_field: str | None = None,
     min_solve_rate: float | None = None,
     max_solve_rate: float | None = None,
+    system_prompt: str | None = None,
     **kwargs,
 ) -> vf.Environment:
     import json
@@ -37,5 +38,9 @@ def correct_answer_reward_func(completion, info, **kwargs) -> float:
         weights=[1.0],
     )
 
-    vf_env = vf.SingleTurnEnv(dataset=train_dataset, rubric=rubric)
+    vf_env = (
+        vf.SingleTurnEnv(dataset=train_dataset, rubric=rubric, system_prompt=system_prompt)
+        if system_prompt is not None
+        else vf.SingleTurnEnv(dataset=train_dataset, rubric=rubric)
+    )
     return vf_env

environments/vf_pydantic_adherence/vf_pydantic_adherence.py

@@ -9,7 +9,7 @@
 from verifiers import Messages, Parser
 
 
-def load_environment() -> vf.Environment:
+def load_environment(system_prompt: str | None = None) -> vf.Environment:
     """
     Loads a custom environment.
     """
@@ -169,10 +169,6 @@ def pydantic_adherence_reward_func(completion, answer, **kwargs):
         weights=[1.0],
     )
 
-    vf_env = vf.SingleTurnEnv(
-        dataset=dataset,
-        parser=parser,
-        rubric=rubric,
-    )
+    vf_env = vf.SingleTurnEnv(dataset=dataset, parser=parser, rubric=rubric, system_prompt=system_prompt) if system_prompt is not None else vf.SingleTurnEnv(dataset=dataset, parser=parser, rubric=rubric)
 
     return vf_env

environments/vf_reverse_text/vf_reverse_text.py

@@ -2,7 +2,7 @@
 from datasets import load_dataset
 
 
-def load_environment() -> vf.Environment:
+def load_environment(system_prompt: str | None = None) -> vf.Environment:
     train_dataset = load_dataset("PrimeIntellect/Reverse-Text-RL", split="train").map(
         lambda x: {
             "question": x["prompt"],
@@ -38,7 +38,8 @@ def lcs_ratio(x: str, y: str) -> float:
         weights=[1.0],
     )
 
-    system_prompt = "Reverse the text character-by-character. Put your answer in <reversed_text> tags."
+    base_system_prompt = "Reverse the text character-by-character. Put your answer in <reversed_text> tags."
+    system_prompt = system_prompt or base_system_prompt
 
     vf_env = vf.SingleTurnEnv(
         dataset=train_dataset,

environments/vf_skywork_math/vf_skywork_math.py

@@ -10,6 +10,7 @@ def load_environment(
     solve_rate_field: str | None = None,
     min_solve_rate: float | None = None,
     max_solve_rate: float | None = None,
+    system_prompt: str | None = None,
     **kwargs,
 ) -> vf.Environment:
     train_dataset = load_dataset("PrimeIntellect/Skywork-OR1-RL-Data-v1-math-prime-rl-format", split="train").map(
@@ -37,5 +38,9 @@ def correct_answer_reward_func(completion, info, **kwargs) -> float:
         weights=[1.0],
     )
 
-    vf_env = vf.SingleTurnEnv(dataset=train_dataset, rubric=rubric)
+    vf_env = (
+        vf.SingleTurnEnv(dataset=train_dataset, rubric=rubric, system_prompt=system_prompt)
+        if system_prompt is not None
+        else vf.SingleTurnEnv(dataset=train_dataset, rubric=rubric)
+    )
     return vf_env

pyproject.toml

@@ -42,6 +42,7 @@ orchestrator = "prime_rl.orchestrator.orchestrator:main"
 inference = "prime_rl.inference.server:main"
 sft = "prime_rl.trainer.sft.train:main"
 eval = "prime_rl.eval.eval:main"
+gepa = "prime_rl.optimizer.gepa.gepa:main"
 
 [build-system]
 requires = ["hatchling"]

src/prime_rl/optimizer/__init__.py

@@ -0,0 +1,8 @@
+"""
+Optimization utilities and algorithms (e.g., GEPA) for prompt and policy search.
+
+This namespace houses optimizers that orchestrate search procedures using the
+existing inference, evaluation, and monitoring infrastructures.
+"""
+
+

src/prime_rl/optimizer/gepa/config.py

@@ -0,0 +1,79 @@
+from pathlib import Path
+from typing import Annotated, Literal
+
+from pydantic import Field
+
+from prime_rl.inference.config import InferenceConfig
+from prime_rl.orchestrator.config import ClientConfig
+from prime_rl.utils.config import LogConfig, ModelConfig, MultiMonitorConfig
+from prime_rl.utils.pydantic_config import BaseConfig, BaseSettings
+
+
+class EvaluateConfig(BaseConfig):
+    """Evaluation settings for GEPA prompt scoring."""
+
+    benchmark: Annotated[str, Field(description="Benchmark/dataset to evaluate on.")] = "math500"
+    rollouts_per_prompt: Annotated[int, Field(ge=1)] = 1
+    pareto_size: Annotated[int, Field(ge=1, description="Number of instances in D_pareto.")] = 64
+    feedback_pool_size: Annotated[int, Field(ge=1, description="Number of instances in D_feedback pool.")] = 256
+    max_tokens: Annotated[int | None, Field(description="Max generated tokens per completion.")] = None
+    min_tokens: Annotated[int, Field(ge=0)] = 0
+
+
+class OperatorsConfig(BaseConfig):
+    """Mutation/crossover operator rates and constraints."""
+
+    mutation_rate: Annotated[float, Field(ge=0, le=1)] = 0.7
+    crossover_rate: Annotated[float, Field(ge=0, le=1)] = 0.3
+    max_prompt_chars: Annotated[int, Field(ge=1)] = 4000
+    enforce_diversity: Annotated[bool, Field(description="Avoid near-duplicate prompts via distance checks.")] = True
+    min_levenshtein_distance: Annotated[int, Field(ge=0)] = 40
+
+
+class SelectionConfig(BaseConfig):
+    """Selection policy settings."""
+
+    strategy: Annotated[Literal["top-k", "tournament"], Field()] = "top-k"
+    k: Annotated[int, Field(ge=1)] = 8
+    tournament_size: Annotated[int, Field(ge=1)] = 3
+    keep_elite: Annotated[int, Field(ge=0)] = 2
+
+
+class GEPAConfig(BaseSettings):
+    """Top-level configuration for the GEPA optimizer."""
+
+    # Model/Client
+    model: ModelConfig = ModelConfig()
+
+    # Evaluation
+    evaluate: EvaluateConfig = EvaluateConfig()
+
+    # Evolution
+    population_size: Annotated[int, Field(ge=2)] = 16
+    generations: Annotated[int, Field(ge=1)] = 10
+    seed: Annotated[int | None, Field(description="Random seed for reproducibility.")] = 42
+
+    # Operators & selection
+    operators: OperatorsConfig = OperatorsConfig()
+    selection: SelectionConfig = SelectionConfig()
+
+    # Budget & minibatch
+    budget_rollouts: Annotated[int, Field(ge=1, description="Total rollout budget for evolution.")] = 200
+    minibatch_size: Annotated[int, Field(ge=1, description="Minibatch size b for acceptance test.")] = 8
+
+    # Logging/monitoring
+    log: LogConfig = LogConfig()
+    monitor: MultiMonitorConfig = MultiMonitorConfig()
+
+    # IO
+    outputs_dir: Annotated[Path, Field(description="Directory for GEPA outputs.")] = Path("outputs")
+    run_name: Annotated[str | None, Field(description="Optional run name for outputs/W&B.")] = None
+
+    # Dev
+    dry_run: Annotated[bool, Field(description="If True, skip real inference and fabricate scores.")] = False
+
+    # Inference server (optional auto-spawn) and client
+    inference: InferenceConfig | None = None
+    client: ClientConfig = ClientConfig()
+
+