training/curriculum.py

"""
Curriculum Scheduler — Unlocks harder scenarios as the agent improves.

Uses rolling 50-episode windows for solve rate calculation to prevent
premature unlocking from lucky streaks. Per the hackathon guide:
"Curriculum is critical for RL convergence."
"""

from __future__ import annotations

import random
from collections import deque
from typing import Dict, List


class CurriculumScheduler:
    """Manages progressive difficulty unlocking for training.

    Starts training on Level 1 only. Unlocks Level 2 once the agent
    achieves > 80% solve rate over the last 50 Level 1 episodes,
    and Level 3 once Level 2 hits > 80% over the last 50 episodes.

    Uses rolling windows (not all-time stats) to prevent premature
    unlocking from a lucky streak.

    Usage:
        scheduler = CurriculumScheduler()
        levels = scheduler.get_active_levels()  # [1]
        scheduler.record_episode(level=1, solved=True)
        # ... after 50+ episodes with >80% solve rate...
        levels = scheduler.get_active_levels()  # [1, 2]
    """

    def __init__(
        self,
        unlock_threshold: float = 0.8,
        window_size: int = 50,
    ) -> None:
        """Initialize the curriculum scheduler.

        Args:
            unlock_threshold: Solve rate threshold to unlock next level (0-1).
            window_size: Number of recent episodes to consider for unlock decisions.
        """
        self.unlock_threshold = unlock_threshold
        self.window_size = window_size

        # Rolling windows per level — stores True/False for solved/failed
        self._windows: Dict[int, deque] = {
            1: deque(maxlen=window_size),
            2: deque(maxlen=window_size),
            3: deque(maxlen=window_size),
        }

        self._unlocked: Dict[int, bool] = {
            1: True,
            2: False,
            3: False,
        }

        self._total_episodes: Dict[int, int] = {1: 0, 2: 0, 3: 0}

    def record_episode(self, level: int, solved: bool) -> None:
        """Record the outcome of an episode for curriculum tracking.

        Args:
            level: The difficulty level of the episode.
            solved: Whether the scenario was solved.
        """
        if level not in self._windows:
            return

        self._windows[level].append(solved)
        self._total_episodes[level] += 1
        self._check_unlocks()

    def get_active_levels(self) -> List[int]:
        """Get currently unlocked difficulty levels.

        Returns:
            List of unlocked level numbers (e.g., [1] or [1, 2]).
        """
        return sorted(lvl for lvl, unlocked in self._unlocked.items() if unlocked)

    def get_window_solve_rate(self, level: int) -> float:
        """Get the solve rate over the rolling window for a level.

        Args:
            level: The level to query.

        Returns:
            Solve rate (0.0-1.0) over the recent window, or 0.0 if no data.
        """
        window = self._windows.get(level, deque())
        if not window:
            return 0.0
        return sum(window) / len(window)

    def _check_unlocks(self) -> None:
        """Check and perform level unlocks based on rolling window stats."""
        # Level 2: unlocks when Level 1 window solve rate >= threshold
        if not self._unlocked[2]:
            window = self._windows[1]
            if len(window) >= self.window_size:
                rate = sum(window) / len(window)
                if rate >= self.unlock_threshold:
                    self._unlocked[2] = True
                    print(f"[Curriculum] ⬆ Level 2 UNLOCKED! "
                          f"(L1 window solve rate: {rate:.1%} over {len(window)} episodes)")

        # Level 3: unlocks when Level 2 window solve rate >= threshold
        if self._unlocked[2] and not self._unlocked[3]:
            window = self._windows[2]
            if len(window) >= self.window_size:
                rate = sum(window) / len(window)
                if rate >= self.unlock_threshold:
                    self._unlocked[3] = True
                    print(f"[Curriculum] ⬆ Level 3 UNLOCKED! "
                          f"(L2 window solve rate: {rate:.1%} over {len(window)} episodes)")

    def sample_level(self) -> int:
        """Sample a level from currently active levels.

        Weighted toward the hardest unlocked level (exactly 50% newest,
        remaining 50% split across other unlocked levels) to keep the
        agent challenged at the frontier.

        Returns:
            A level number to train on.
        """
        active = self.get_active_levels()
        if len(active) == 1:
            return active[0]

        # 50% chance for the hardest unlocked level
        newest = max(active)
        if random.random() < 0.5:
            return newest

        others = [lvl for lvl in active if lvl != newest]
        return random.choice(others) if others else newest

    def update_stats(self, level: int, solve_rate: float, episodes: int = 0) -> None:
        """Bulk-update stats from replay buffer (for API/UI compatibility).

        This is a convenience method that doesn't use rolling windows —
        prefer record_episode() for training accuracy.

        Args:
            level: The level to update.
            solve_rate: Current solve rate for this level.
            episodes: Total episodes at this level.
        """
        if level not in self._total_episodes:
            return
        if episodes > 0:
            self._total_episodes[level] = episodes

    def get_status(self) -> Dict:
        """Get the current curriculum status.

        Returns:
            Dict with level stats, unlock status, and window solve rates.
        """
        return {
            level: {
                "unlocked": self._unlocked[level],
                "window_solve_rate": round(self.get_window_solve_rate(level), 3),
                "window_size": len(self._windows[level]),
                "total_episodes": self._total_episodes[level],
            }
            for level in [1, 2, 3]
        }

    def force_unlock(self, level: int) -> None:
        """Force-unlock a level (for debugging/testing).

        Args:
            level: Level to unlock.
        """
        if level in self._unlocked:
            self._unlocked[level] = True