Add auditing metadata to RL scheduler outputs

court_scheduler_rl.py

@@ -97,13 +97,21 @@ class InteractivePipeline:
         console.print("\n[bold cyan]Step 1/7: EDA & Parameter Extraction[/bold cyan]")
         
         # Check if EDA was run recently
+        from src import eda_config
+
         param_dir = Path("reports/figures").glob("v0.4.0_*/params")
-        recent_params = any(p.exists() and 
+        recent_params = any(p.exists() and
                           (datetime.now() - datetime.fromtimestamp(p.stat().st_mtime)).days < 1
                           for p in param_dir)
-        
+
         if recent_params and not Confirm.ask("EDA parameters found. Regenerate?", default=False):
             console.print("  [green]OK[/green] Using existing EDA parameters")
+            self.output.record_eda_metadata(
+                version=eda_config.VERSION,
+                used_cached=True,
+                params_path=self.output.eda_params,
+                figures_path=self.output.eda_figures,
+            )
             return
             
         with Progress(
@@ -127,10 +135,16 @@ class InteractivePipeline:
             run_load_and_clean()
             run_exploration()
             run_parameter_export()
-            
+
             progress.update(task, completed=True)
-        
+
         console.print("  [green]OK[/green] EDA pipeline complete")
+        self.output.record_eda_metadata(
+            version=eda_config.VERSION,
+            used_cached=False,
+            params_path=self.output.eda_params,
+            figures_path=self.output.eda_figures,
+        )
     
     def _step_2_data_generation(self):
         """Step 2: Generate Training Data"""
@@ -169,7 +183,10 @@ class InteractivePipeline:
         console.print(f"  Episodes: {self.config.rl_training.episodes}, Learning Rate: {self.config.rl_training.learning_rate}")
         
         model_file = self.output.trained_model_file
-        
+
+        def _safe_mean(values: List[float]) -> float:
+            return sum(values) / len(values) if values else 0.0
+
         with Progress(
             SpinnerColumn(),
             TextColumn("[progress.description]{task.description}"),
@@ -201,12 +218,63 @@ class InteractivePipeline:
                 episode_length=rl_cfg.episode_length_days,
                 verbose=False  # Disable internal printing
             )
-            
+
             progress.update(training_task, completed=rl_cfg.episodes)
-            
+
             # Save trained agent
             agent.save(model_file)
-            
+
+            # Persist training stats for downstream consumers
+            self.output.save_training_stats(training_stats)
+
+            # Run a lightweight evaluation sweep for summary metrics
+            evaluation_stats = None
+            try:
+                from rl.training import evaluate_agent
+                from scheduler.data.case_generator import CaseGenerator
+
+                eval_gen = CaseGenerator(
+                    start=date.today(),
+                    end=date.today() + timedelta(days=60),
+                    seed=self.config.seed + 99,
+                )
+                eval_cases = eval_gen.generate(min(rl_cfg.cases_per_episode, 500), stage_mix_auto=True)
+                evaluation_stats = evaluate_agent(
+                    agent=agent,
+                    test_cases=eval_cases,
+                    episodes=5,
+                    episode_length=rl_cfg.episode_length_days,
+                )
+                self.output.save_evaluation_stats(evaluation_stats)
+            except Exception as eval_err:
+                console.print(f"  [yellow]WARNING[/yellow] Evaluation skipped: {eval_err}")
+
+            training_summary = {
+                "episodes": rl_cfg.episodes,
+                "cases_per_episode": rl_cfg.cases_per_episode,
+                "episode_length_days": rl_cfg.episode_length_days,
+                "learning_rate": rl_cfg.learning_rate,
+                "epsilon": {
+                    "initial": rl_cfg.initial_epsilon,
+                    "final": agent.epsilon,
+                },
+                "reward": {
+                    "mean": _safe_mean(training_stats.get("total_rewards", [])),
+                    "final": training_stats.get("total_rewards", [0])[-1] if training_stats.get("total_rewards") else 0.0,
+                },
+                "disposal_rate": {
+                    "mean": _safe_mean(training_stats.get("disposal_rates", [])),
+                    "final": training_stats.get("disposal_rates", [0])[-1] if training_stats.get("disposal_rates") else 0.0,
+                },
+                "states_explored_final": training_stats.get("states_explored", [len(agent.q_table)])[-1]
+                if training_stats.get("states_explored")
+                else len(agent.q_table),
+                "q_table_size": len(agent.q_table),
+                "total_updates": getattr(agent, "total_updates", 0),
+            }
+
+            self.output.record_training_summary(training_summary, evaluation_stats)
+
             # Create symlink in models/ for backwards compatibility
             self.output.create_model_symlink()
         
@@ -270,18 +338,38 @@ class InteractivePipeline:
                 
                 sim = CourtSim(cfg, policy_cases)
                 result = sim.run()
-                
+
                 progress.update(task, completed=100)
-                
+
                 results[policy] = {
                     'result': result,
                     'cases': policy_cases,  # Use the deep-copied cases for this simulation
                     'sim': sim,
                     'dir': policy_dir
                 }
-            
+
             console.print(f"    [green]OK[/green] {result.disposals:,} disposals ({result.disposals/len(cases):.1%})")
-        
+
+            allocator_stats = sim.allocator.get_utilization_stats()
+            backlog = sum(1 for c in policy_cases if not c.is_disposed)
+
+            kpis = {
+                "policy": policy,
+                "disposals": result.disposals,
+                "disposal_rate": result.disposals / len(policy_cases),
+                "utilization": result.utilization,
+                "hearings_total": result.hearings_total,
+                "hearings_heard": result.hearings_heard,
+                "hearings_adjourned": result.hearings_adjourned,
+                "backlog": backlog,
+                "backlog_rate": backlog / len(policy_cases) if policy_cases else 0,
+                "fairness_gini": allocator_stats.get("load_balance_gini"),
+                "avg_daily_load": allocator_stats.get("avg_daily_load"),
+                "capacity_rejections": allocator_stats.get("capacity_rejections"),
+            }
+
+            self.output.record_simulation_kpis(policy, kpis)
+
         self.sim_results = results
         console.print(f"  [green]OK[/green] All simulations complete")
     

scheduler/utils/output_manager.py

@@ -6,7 +6,7 @@ No scattered files, no duplicate saves, single source of truth per run.
 
 from pathlib import Path
 from datetime import datetime
-from typing import Optional
+from typing import Optional, Dict, Any
 import json
 from dataclasses import asdict
 
@@ -30,7 +30,8 @@ class OutputManager:
             base_dir: Base directory for all outputs (default: outputs/runs)
         """
         self.run_id = run_id or f"run_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
-        
+        self.created_at = datetime.now().isoformat()
+
         # Base paths
         project_root = Path(__file__).parent.parent.parent
         self.base_dir = base_dir or (project_root / "outputs" / "runs")
@@ -49,6 +50,9 @@ class OutputManager:
         
         # Reports subdirectories
         self.visualizations_dir = self.reports_dir / "visualizations"
+
+        # Metadata paths
+        self.run_record_file = self.run_dir / "run_record.json"
         
     def create_structure(self):
         """Create all output directories."""
@@ -64,10 +68,18 @@ class OutputManager:
             self.visualizations_dir,
         ]:
             dir_path.mkdir(parents=True, exist_ok=True)
-    
+
+        # Initialize run record with creation metadata if missing
+        if not self.run_record_file.exists():
+            self._update_run_record("run", {
+                "run_id": self.run_id,
+                "created_at": self.created_at,
+                "base_dir": str(self.run_dir),
+            })
+
     def save_config(self, config):
         """Save pipeline configuration to run directory.
-        
+
         Args:
             config: PipelineConfig or any dataclass
         """
@@ -76,6 +88,45 @@ class OutputManager:
             # Handle nested dataclasses (like rl_training)
             config_dict = asdict(config) if hasattr(config, '__dataclass_fields__') else config
             json.dump(config_dict, f, indent=2, default=str)
+
+        self._update_run_record("config", {
+            "path": str(config_path),
+            "timestamp": datetime.now().isoformat(),
+        })
+
+    def save_training_stats(self, training_stats: Dict[str, Any]):
+        """Persist raw training statistics for auditing and dashboards."""
+
+        self.training_dir.mkdir(parents=True, exist_ok=True)
+        with open(self.training_stats_file, "w", encoding="utf-8") as f:
+            json.dump(training_stats, f, indent=2, default=str)
+
+    def save_evaluation_stats(self, evaluation_stats: Dict[str, Any]):
+        """Persist evaluation metrics for downstream analysis."""
+
+        eval_path = self.training_dir / "evaluation.json"
+        with open(eval_path, "w", encoding="utf-8") as f:
+            json.dump(evaluation_stats, f, indent=2, default=str)
+
+        self._update_run_record("evaluation", {
+            "path": str(eval_path),
+            "timestamp": datetime.now().isoformat(),
+        })
+
+    def record_training_summary(self, summary: Dict[str, Any], evaluation: Optional[Dict[str, Any]] = None):
+        """Save aggregated training/evaluation summary for dashboards."""
+
+        summary_path = self.training_dir / "summary.json"
+        payload = {
+            "summary": summary,
+            "evaluation": evaluation,
+            "updated_at": datetime.now().isoformat(),
+        }
+
+        with open(summary_path, "w", encoding="utf-8") as f:
+            json.dump(payload, f, indent=2, default=str)
+
+        self._update_run_record("training", payload)
     
     def get_policy_dir(self, policy_name: str) -> Path:
         """Get simulation directory for a specific policy.
@@ -102,7 +153,37 @@ class OutputManager:
         cause_list_dir = self.get_policy_dir(policy_name) / "cause_lists"
         cause_list_dir.mkdir(parents=True, exist_ok=True)
         return cause_list_dir
-    
+
+    def record_eda_metadata(self, version: str, used_cached: bool, params_path: Path, figures_path: Path):
+        """Record EDA version/timestamp for auditability."""
+
+        payload = {
+            "version": version,
+            "timestamp": datetime.now().isoformat(),
+            "used_cached": used_cached,
+            "params_path": str(params_path),
+            "figures_path": str(figures_path),
+        }
+
+        self._update_run_record("eda", payload)
+
+    def record_simulation_kpis(self, policy: str, kpis: Dict[str, Any]):
+        """Persist simulation KPIs per policy for dashboards."""
+
+        policy_dir = self.get_policy_dir(policy)
+        metrics_path = policy_dir / "metrics.json"
+        with open(metrics_path, "w", encoding="utf-8") as f:
+            json.dump(kpis, f, indent=2, default=str)
+
+        record = self._load_run_record()
+        simulation_section = record.get("simulation", {})
+        simulation_section[policy] = kpis
+        record["simulation"] = simulation_section
+        record["updated_at"] = datetime.now().isoformat()
+
+        with open(self.run_record_file, "w", encoding="utf-8") as f:
+            json.dump(record, f, indent=2, default=str)
+
     @property
     def training_cases_file(self) -> Path:
         """Path to generated training cases CSV."""
@@ -152,9 +233,38 @@ class OutputManager:
             # Fallback: copy file if symlinks not supported (Windows without dev mode)
             import shutil
             shutil.copy2(target, symlink_path)
-    
+
     def __str__(self) -> str:
         return f"OutputManager(run_id='{self.run_id}', run_dir='{self.run_dir}')"
-    
+
     def __repr__(self) -> str:
         return self.__str__()
+
+    # ------------------------------------------------------------------
+    # Internal helpers
+    # ------------------------------------------------------------------
+    def _load_run_record(self) -> Dict[str, Any]:
+        """Load run record JSON, providing defaults if missing."""
+
+        if self.run_record_file.exists():
+            try:
+                with open(self.run_record_file, "r", encoding="utf-8") as f:
+                    return json.load(f)
+            except json.JSONDecodeError:
+                pass
+
+        return {
+            "run_id": self.run_id,
+            "created_at": self.created_at,
+        }
+
+    def _update_run_record(self, section: str, payload: Dict[str, Any]):
+        """Upsert a section within the consolidated run record."""
+
+        record = self._load_run_record()
+        record.setdefault("sections", {})
+        record["sections"][section] = payload
+        record["updated_at"] = datetime.now().isoformat()
+
+        with open(self.run_record_file, "w", encoding="utf-8") as f:
+            json.dump(record, f, indent=2, default=str)