Merge pull request #261 from totoluo/feat-export-trace

feature: Add evolution trace logging for RL training support

Author: codelion

configs/default_config.yaml

@@ -159,3 +159,14 @@ evaluator:
   # LLM-based feedback (experimental)
   use_llm_feedback: false             # Use LLM to evaluate code quality
   llm_feedback_weight: 0.1            # Weight for LLM feedback in final score
+
+# Evolution trace configuration
+# Logs detailed traces of program evolution for RL training and analysis
+evolution_trace:
+  enabled: false                      # Enable evolution trace logging
+  format: 'jsonl'                     # Output format: 'jsonl', 'json', or 'hdf5'
+  include_code: false                 # Include full program code in traces
+  include_prompts: true               # Include prompts and LLM responses
+  output_path: null                   # Path for trace output (defaults to output_dir/evolution_trace.{format})
+  buffer_size: 10                     # Number of traces to buffer before writing
+  compress: false                     # Compress output file (jsonl only)

openevolve/config.py

@@ -313,6 +313,19 @@ class EvaluatorConfig:
     max_artifact_storage: int = 100 * 1024 * 1024  # 100MB per program
 
 
+@dataclass
+class EvolutionTraceConfig:
+    """Configuration for evolution trace logging"""
+    
+    enabled: bool = False
+    format: str = "jsonl"  # Options: "jsonl", "json", "hdf5"
+    include_code: bool = False
+    include_prompts: bool = True
+    output_path: Optional[str] = None
+    buffer_size: int = 10
+    compress: bool = False
+
+
 @dataclass
 class Config:
     """Master configuration for OpenEvolve"""
@@ -330,6 +343,7 @@ class Config:
     prompt: PromptConfig = field(default_factory=PromptConfig)
     database: DatabaseConfig = field(default_factory=DatabaseConfig)
     evaluator: EvaluatorConfig = field(default_factory=EvaluatorConfig)
+    evolution_trace: EvolutionTraceConfig = field(default_factory=EvolutionTraceConfig)
 
     # Evolution settings
     diff_based_evolution: bool = True
@@ -355,7 +369,7 @@ def from_dict(cls, config_dict: Dict[str, Any]) -> "Config":
 
         # Update top-level fields
         for key, value in config_dict.items():
-            if key not in ["llm", "prompt", "database", "evaluator"] and hasattr(config, key):
+            if key not in ["llm", "prompt", "database", "evaluator", "evolution_trace"] and hasattr(config, key):
                 setattr(config, key, value)
 
         # Update nested configs
@@ -378,6 +392,8 @@ def from_dict(cls, config_dict: Dict[str, Any]) -> "Config":
             config.database.random_seed = config.random_seed
         if "evaluator" in config_dict:
             config.evaluator = EvaluatorConfig(**config_dict["evaluator"])
+        if "evolution_trace" in config_dict:
+            config.evolution_trace = EvolutionTraceConfig(**config_dict["evolution_trace"])
 
         return config
 
@@ -446,6 +462,15 @@ def to_dict(self) -> Dict[str, Any]:
                 "use_llm_feedback": self.evaluator.use_llm_feedback,
                 "llm_feedback_weight": self.evaluator.llm_feedback_weight,
             },
+            "evolution_trace": {
+                "enabled": self.evolution_trace.enabled,
+                "format": self.evolution_trace.format,
+                "include_code": self.evolution_trace.include_code,
+                "include_prompts": self.evolution_trace.include_prompts,
+                "output_path": self.evolution_trace.output_path,
+                "buffer_size": self.evolution_trace.buffer_size,
+                "compress": self.evolution_trace.compress,
+            },
             # Evolution settings
             "diff_based_evolution": self.diff_based_evolution,
             "max_code_length": self.max_code_length,

openevolve/controller.py

@@ -14,6 +14,7 @@
 from openevolve.config import Config, load_config
 from openevolve.database import Program, ProgramDatabase
 from openevolve.evaluator import Evaluator
+from openevolve.evolution_trace import EvolutionTracer
 from openevolve.llm.ensemble import LLMEnsemble
 from openevolve.prompt.sampler import PromptSampler
 from openevolve.process_parallel import ProcessParallelController
@@ -163,6 +164,29 @@ def __init__(
 
         logger.info(f"Initialized OpenEvolve with {initial_program_path}")
 
+        # Initialize evolution tracer
+        if self.config.evolution_trace.enabled:
+            trace_output_path = self.config.evolution_trace.output_path
+            if not trace_output_path:
+                # Default to output_dir/evolution_trace.{format}
+                trace_output_path = os.path.join(
+                    self.output_dir, 
+                    f"evolution_trace.{self.config.evolution_trace.format}"
+                )
+            
+            self.evolution_tracer = EvolutionTracer(
+                output_path=trace_output_path,
+                format=self.config.evolution_trace.format,
+                include_code=self.config.evolution_trace.include_code,
+                include_prompts=self.config.evolution_trace.include_prompts,
+                enabled=True,
+                buffer_size=self.config.evolution_trace.buffer_size,
+                compress=self.config.evolution_trace.compress
+            )
+            logger.info(f"Evolution tracing enabled: {trace_output_path}")
+        else:
+            self.evolution_tracer = None
+
         # Initialize improved parallel processing components
         self.parallel_controller = None
 
@@ -276,7 +300,7 @@ async def run(
         # Initialize improved parallel processing
         try:
             self.parallel_controller = ProcessParallelController(
-                self.config, self.evaluation_file, self.database
+                self.config, self.evaluation_file, self.database, self.evolution_tracer
             )
 
             # Set up signal handlers for graceful shutdown
@@ -319,6 +343,11 @@ def force_exit_handler(signum, frame):
             if self.parallel_controller:
                 self.parallel_controller.stop()
                 self.parallel_controller = None
+            
+            # Close evolution tracer
+            if self.evolution_tracer:
+                self.evolution_tracer.close()
+                logger.info("Evolution tracer closed")
 
         # Get the best program
         best_program = None