Kalletlamadhav/sql-optimization-env

1# curriculum/progress_tracker.py2"""3ProgressTracker4---------------5Records per-agent episode history so you can query an agent's score6trajectory independently of the shared CurriculumEngine state.7 8Stored as a JSON file at /tmp/progress_<agent_id>.json so it survives9container restarts within the same session.10 11Usage:12    tracker = ProgressTracker(agent_id='baseline_gpt4o')13    tracker.record(task_id='gst_missing_index', reward=0.78, level=2)14    print(tracker.mean_reward(last_n=10))15"""16 17import json18from pathlib import Path19from typing import Optional20 21 22class ProgressTracker:23    def __init__(self, agent_id: str = 'default'):24        self.agent_id = agent_id25        self._path = Path(f'/tmp/progress_{agent_id}.json')26        self._history: list[dict] = []27        self._load()28 29    # ── Public API ────────────────────────────────────────────────────────────30 31    def record(self, task_id: str, reward: float, level: int,32               speedup: float = 0.0, hack: Optional[str] = None) -> None:33        """Append one episode result."""34        entry = {35            'task_id':  task_id,36            'reward':   round(reward, 4),37            'level':    level,38            'speedup':  round(speedup, 2),39            'hack':     hack,40            'episode':  len(self._history) + 1,41        }42        self._history.append(entry)43        self._save()44 45    def mean_reward(self, last_n: int = 0) -> float:46        """Average reward over the last *last_n* episodes (0 = all)."""47        data = self._history[-last_n:] if last_n else self._history48        if not data:49            return 0.050        return round(sum(e['reward'] for e in data) / len(data), 4)51 52    def best_reward(self, task_id: Optional[str] = None) -> float:53        """Best reward ever seen, optionally filtered by task."""54        data = [e for e in self._history if not task_id or e['task_id'] == task_id]55        return max((e['reward'] for e in data), default=0.0)56 57    def total_episodes(self) -> int:58        return len(self._history)59 60    def task_summary(self) -> dict:61        """Return {task_id: {'count': N, 'mean': X, 'best': Y}} for all tasks."""62        tasks: dict[str, list[float]] = {}63        for e in self._history:64            tasks.setdefault(e['task_id'], []).append(e['reward'])65        return {66            tid: {67                'count': len(rewards),68                'mean':  round(sum(rewards) / len(rewards), 4),69                'best':  max(rewards),70            }71            for tid, rewards in tasks.items()72        }73 74    def reset(self) -> None:75        """Clear all history for this agent."""76        self._history = []77        if self._path.exists():78            self._path.unlink()79 80    # ── Persistence ───────────────────────────────────────────────────────────81 82    def _load(self) -> None:83        if self._path.exists():84            try:85                with open(self._path) as f:86                    self._history = json.load(f)87            except (json.JSONDecodeError, OSError):88                self._history = []89 90    def _save(self) -> None:91        try:92            with open(self._path, 'w') as f:93                json.dump(self._history, f)94        except OSError:95            pass  # Non-fatal — state lost on restart, not a correctness issue