""" Replay Buffer — SQLAlchemy-backed episode storage and sampling. Stores every episode as structured data in SQLite, supports batch sampling for training, and scenario-based querying for analysis. """ from __future__ import annotations from contextlib import contextmanager import json import uuid import random from datetime import datetime, timezone from typing import Dict, Iterator, List, Optional from sqlalchemy import func, desc from sqlalchemy.orm import Session from replay.models import EpisodeRecord, StepRecord, create_database class ReplayBuffer: """Persistent replay buffer backed by SQLite. Stores episodes with their steps, supports batch sampling for GRPO training, and provides analytics queries. Usage: buffer = ReplayBuffer("sqlite:///episodes.db") buffer.store_episode(episode_data) batch = buffer.sample_batch(32) scenarios = buffer.get_by_scenario("missing_flask") """ def __init__(self, db_url: str = "sqlite:///replay_buffer.db") -> None: """Initialize the replay buffer. Args: db_url: SQLAlchemy database URL for episode storage. """ self.db_url = db_url self._session_factory = create_database(db_url) @contextmanager def _get_session(self) -> Iterator[Session]: """Yield a managed database session and always close it.""" session = self._session_factory() try: yield session finally: session.close() def store_episode( self, scenario_id: str, level: int, steps: List[Dict], total_reward: float, solved: bool, training_episode: int | None = None, ) -> str: """Store a complete episode in the buffer. Args: scenario_id: Which scenario was attempted. level: Difficulty level. steps: List of step dicts with observation, action, result, reward. total_reward: Sum of all step rewards. solved: Whether the scenario was solved. training_episode: Training episode number, if applicable. Returns: The generated episode UUID. """ episode_id = str(uuid.uuid4()) with self._get_session() as session: episode = EpisodeRecord( episode_id=episode_id, scenario_id=scenario_id, level=level, total_reward=total_reward, solved=solved, total_steps=len(steps), timestamp=datetime.now(timezone.utc), training_episode=training_episode, ) session.add(episode) for step_data in steps: step = StepRecord( episode=episode, step_number=step_data.get("step", 0), observation_json=json.dumps(step_data.get("observation", {})), action=step_data.get("action", ""), result_json=json.dumps(step_data.get("result", {})), reward=step_data.get("reward", 0.0), reward_breakdown_json=json.dumps(step_data.get("reward_breakdown", {})), error_type=step_data.get("error_type", "unknown"), ) session.add(step) session.commit() return episode_id def sample_batch(self, n: int, level: int | None = None) -> List[Dict]: """Sample a random batch of episodes for training. Args: n: Number of episodes to sample. level: If provided, only sample from this level. Returns: List of episode dicts with full step details. """ with self._get_session() as session: query = session.query(EpisodeRecord) if level is not None: query = query.filter(EpisodeRecord.level == level) total = query.count() if total == 0: return [] # Random offset sampling if total <= n: episodes = query.all() else: # Get random IDs all_ids = [r.id for r in query.with_entities(EpisodeRecord.id).all()] sampled_ids = random.sample(all_ids, min(n, len(all_ids))) episodes = query.filter(EpisodeRecord.id.in_(sampled_ids)).all() return [ep.to_dict() for ep in episodes] def get_by_scenario(self, scenario_id: str, limit: int = 100) -> List[Dict]: """Get all episodes for a specific scenario. Args: scenario_id: Scenario ID to filter by. limit: Maximum number of episodes to return. Returns: List of episode dicts, newest first. """ with self._get_session() as session: episodes = ( session.query(EpisodeRecord) .filter(EpisodeRecord.scenario_id == scenario_id) .order_by(desc(EpisodeRecord.timestamp)) .limit(limit) .all() ) return [ep.to_dict() for ep in episodes] def get_episode(self, episode_id: str) -> Optional[Dict]: """Get a specific episode by its UUID. Args: episode_id: The episode UUID string. Returns: Episode dict or None if not found. """ with self._get_session() as session: episode = ( session.query(EpisodeRecord) .filter(EpisodeRecord.episode_id == episode_id) .first() ) if episode: return episode.to_dict() return None def get_stats(self) -> Dict: """Get aggregate statistics across all stored episodes. Returns: Dict with solve rates, mean rewards, counts per level. """ with self._get_session() as session: stats: Dict = {"total_episodes": 0, "levels": {}} total = session.query(func.count(EpisodeRecord.id)).scalar() stats["total_episodes"] = total or 0 for level in [1, 2, 3]: level_query = session.query(EpisodeRecord).filter(EpisodeRecord.level == level) level_count = level_query.count() if level_count == 0: stats["levels"][level] = { "count": 0, "solve_rate": 0.0, "mean_reward": 0.0, "mean_steps": 0.0, } continue solved_count = level_query.filter(EpisodeRecord.solved == True).count() mean_reward = session.query( func.avg(EpisodeRecord.total_reward) ).filter(EpisodeRecord.level == level).scalar() or 0.0 mean_steps = session.query( func.avg(EpisodeRecord.total_steps) ).filter(EpisodeRecord.level == level).scalar() or 0.0 stats["levels"][level] = { "count": level_count, "solve_rate": solved_count / level_count if level_count > 0 else 0.0, "mean_reward": round(float(mean_reward), 2), "mean_steps": round(float(mean_steps), 2), } # Per-scenario stats scenario_stats = {} scenarios = session.query( EpisodeRecord.scenario_id ).distinct().all() for (sid,) in scenarios: sc_query = session.query(EpisodeRecord).filter(EpisodeRecord.scenario_id == sid) sc_count = sc_query.count() sc_solved = sc_query.filter(EpisodeRecord.solved == True).count() scenario_stats[sid] = { "count": sc_count, "solve_rate": sc_solved / sc_count if sc_count > 0 else 0.0, } stats["scenarios"] = scenario_stats return stats def get_recent(self, n: int = 20) -> List[Dict]: """Get the most recent episodes. Args: n: Number of episodes to return. Returns: List of episode dicts, newest first. """ with self._get_session() as session: episodes = ( session.query(EpisodeRecord) .order_by(desc(EpisodeRecord.timestamp)) .limit(n) .all() ) return [ep.to_dict() for ep in episodes] @property def size(self) -> int: """Total number of episodes in the buffer.""" with self._get_session() as session: return session.query(func.count(EpisodeRecord.id)).scalar() or 0