experiment.py

import functools
import json
import tempfile
import time
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional

import mlflow
from omegaconf import DictConfig, OmegaConf
from tqdm import tqdm

from agents.Agent import Agent
from poolflip.flipit_v3 import PoolFlipEnv
from poolflip.wrappers.envs.ExperimentSaver import ExperimentSaverParallelWrapper


def handle_experiment_errors(
    raise_on_error: bool = True,
    log_error: bool = True,
    error_callback: Optional[
        Callable[[Exception, DictConfig, Dict[str, Any]], None]
    ] = None,
):
    """
    Decorator to handle errors in experiment execution.

    Args:
        raise_on_error: Whether to re-raise the exception after handling
        log_error: Whether to print error information
        error_callback: Optional callback function to handle errors
    """

    def decorator(func):
        @functools.wraps(func)
        def wrapper(cfg: DictConfig, agents: Dict[str, Any], *args, **kwargs):
            try:
                return func(cfg, agents, *args, **kwargs)
            except Exception as e:
                if log_error:
                    print(
                        f"Error running experiment: {e} with config: {cfg} and agents: {agents}"
                    )
                if error_callback:
                    error_callback(e, cfg, agents)
                if raise_on_error:
                    raise e

        return wrapper

    return decorator


class EarlyStopping:
    """A class to handle early stopping based on reward convergence.

    This class implements early stopping logic that can be used to stop training
    when the reward improvement plateaus. Uses non-overlapping windows (bins) to
    detect convergence in RL rewards.
    """

    def __init__(
        self,
        patience: int = 10,
        min_delta: float = 0.01,
        enabled: bool = True,
        window_size: int = 5,
    ):
        """Initialize the early stopping module.

        Args:
            patience: Number of consecutive bins that must be within min_delta
            min_delta: Maximum allowed difference between bin averages
            enabled: Whether early stopping is enabled
            window_size: Number of episodes per bin
        """
        self.patience = patience
        self.min_delta = min_delta
        self.enabled = enabled
        self.window_size = window_size
        self.reward_history = []
        self.bin_averages = []
        self.triggered = False
        self.episode_triggered = -1

    def _calculate_bin_average(self, start_idx: int) -> float:
        """Calculate the average reward for a bin.

        Args:
            start_idx: Starting index for the bin

        Returns:
            The average reward for the bin
        """
        end_idx = min(start_idx + self.window_size, len(self.reward_history))
        if end_idx - start_idx < self.window_size:
            return float("-inf")
        return sum(self.reward_history[start_idx:end_idx]) / self.window_size

    def _check_convergence(self) -> bool:
        """Check if the last patience bins are within min_delta of each other.

        Returns:
            True if the bins are converged, False otherwise
        """
        if len(self.bin_averages) < self.patience:
            return False

        # Get the last patience bins
        recent_bins = self.bin_averages[-self.patience :]

        # Check if all bins are within min_delta of each other
        min_avg = min(recent_bins)
        max_avg = max(recent_bins)
        return (max_avg - min_avg) <= self.min_delta

    def check(self, episode_rewards: Dict[str, float], episode_ndx: int) -> bool:
        """Check if training should be stopped based on reward convergence.

        Args:
            episode_rewards: Dictionary mapping agent names to their episode rewards
            episode_ndx: Current episode index

        Returns:
            True if training should be stopped, False otherwise
        """
        if not self.enabled:
            return False

        # Calculate average reward across all agents for this episode
        avg_reward = sum(episode_rewards.values()) / len(episode_rewards)
        self.reward_history.append(avg_reward)

        # Check if we have enough episodes for a new bin
        if len(self.reward_history) % self.window_size == 0:
            bin_start = len(self.reward_history) - self.window_size
            bin_avg = self._calculate_bin_average(bin_start)
            if bin_avg != float("-inf"):
                self.bin_averages.append(bin_avg)

                # Check for convergence
                if self._check_convergence():
                    self.triggered = True
                    self.episode_triggered = episode_ndx
                    print(
                        f"Early stopping triggered after {len(self.bin_averages)} "
                        f"bins. Last {self.patience} bins converged within "
                        f"{self.min_delta:.3f}. Bin averages: "
                        f"{self.bin_averages[-self.patience :]}"
                    )
                    return True

        return False

    def reset(self):
        """Reset the early stopping state."""
        self.reward_history = []
        self.bin_averages = []
        self.triggered = False
        self.episode_triggered = -1

    def get_status(self) -> Dict[str, Any]:
        """Get the current status of early stopping.

        Returns:
            Dictionary containing early stopping status and parameters
        """
        return {
            "enabled": self.enabled,
            "patience": self.patience,
            "min_delta": self.min_delta,
            "window_size": self.window_size,
            "triggered": self.triggered,
            "episode_triggered": self.episode_triggered,
            "num_bins": len(self.bin_averages),
            "current_bin_average": self.bin_averages[-1] if self.bin_averages else None,
            "last_patience_bins": (
                self.bin_averages[-self.patience :]
                if len(self.bin_averages) >= self.patience
                else None
            ),
        }


def save_policies(agent_name_to_agent: Dict[str, Agent], output_dir: Path, suffix: str):
    for name, agent in agent_name_to_agent.items():
        agent.save(path=output_dir / f"{name}_{suffix}.pth")
        print(f"Saved {name} policy to {output_dir / f'{name}_{suffix}.pth'}")


class ExperimentRunner:
    def __init__(
        self,
        env: Any,
        save_every_k_episodes: int,
        output_dir: Path,
        early_stopping: Optional[EarlyStopping] = None,
    ):
        """Initialize the experiment runner.

        Args:
            env: The environment to run experiments in
            save_every_k_episodes: Save checkpoints every k episodes
            output_dir: Directory to save outputs
            early_stopping: Optional early stopping module to use
        """
        self.env = env
        self.save_every_k_episodes = save_every_k_episodes
        self.output_dir = output_dir
        self.output_dir.mkdir(parents=True, exist_ok=True)
        self.early_stopping = early_stopping

    def run(self, num_episodes: int, num_steps: int) -> Dict[str, Dict[int, float]]:
        """Run the experiment.

        Args:
            num_episodes: Number of episodes to run
            num_steps: Number of steps per episode

        Returns:
            Dictionary mapping agent names to dictionaries of episode rewards
        """
        print(
            "[DEBUG] Running experiment with num_episodes: ",
            num_episodes,
            " and num_steps: ",
            num_steps,
            "and early_stopping: ",
            self.early_stopping,
            "and agents: ",
            self.env.agent_name_to_agent,
        )
        save_policies(self.env.agent_name_to_agent, self.output_dir, "initial_policy")
        per_episode_rewards = {}
        for episode_ndx in tqdm(range(num_episodes), desc="Running experiment"):
            if episode_ndx not in per_episode_rewards:
                per_episode_rewards[episode_ndx] = {}
            observations, infos = self.env.reset()
            for step_ndx in range(num_steps):
                actions = {
                    agent_name: self.env.agent_name_to_agent[agent_name].get_action(
                        observations[agent_name]
                    )
                    for agent_name in self.env.agents
                }
                observations, rewards, terminations, truncations, infos = self.env.step(
                    actions
                )
                should_break = any(terminations.values()) or any(truncations.values())
                if should_break:
                    break
                for agent_name in self.env.possible_agents:
                    agent = self.env.agent_name_to_agent[agent_name]
                    agent.store(
                        rewards[agent_name],
                        terminations[agent_name] or truncations[agent_name],
                    )
                    if agent_name not in per_episode_rewards[episode_ndx]:
                        per_episode_rewards[episode_ndx][agent_name] = 0
                    per_episode_rewards[episode_ndx][agent_name] += rewards[agent_name]

            # Check for early stopping after each episode
            if self.early_stopping and self.early_stopping.check(
                per_episode_rewards[episode_ndx], episode_ndx
            ):
                print(f"Stopping training at episode {episode_ndx}")
                break

            if episode_ndx % self.save_every_k_episodes == 0:
                for name, agent in self.env.agent_name_to_agent.items():
                    agent.save(
                        path=self.output_dir / f"{name}_{episode_ndx}_policy.pth"
                    )

        # Save final policies
        save_policies(self.env.agent_name_to_agent, self.output_dir, "final_policy")
        self.env.reset()  # Triggers the save for the last episode.

        return per_episode_rewards


def run_experiment_with_config(
    cfg: DictConfig,
    agents: Dict[str, Agent],
    experiment_type: str = "experiment",
    tags: List[str] = None,
) -> Dict[str, Dict[int, float]]:
    """Run an experiment with the given configuration and agents.

    Args:
        cfg: Configuration for the experiment
        agents: Dictionary mapping agent names to agent instances
        experiment_type: Type of experiment being run
        tags: List of tags for the experiment

    Returns:
        Dictionary mapping agent names to dictionaries of episode rewards
    """
    if tags is None:
        tags = ["experiment"]

    with tempfile.TemporaryDirectory() as tmpdir:
        return run_experiment(cfg, agents, Path(tmpdir))


@handle_experiment_errors(raise_on_error=True)
def wrappable_run_experiment(
    cfg: DictConfig, agents: Dict[str, Agent], experiment_type: str, tags: List[str]
) -> Dict[str, Dict[int, float]]:
    """Run an experiment with error handling.

    This function is a thin wrapper around run_experiment_with_config that adds error handling.

    Args:
        cfg: Configuration for the experiment
        agents: Dictionary mapping agent names to agent instances
        experiment_type: Type of experiment being run
        tags: List of tags for the experiment

    Returns:
        Dictionary mapping agent names to dictionaries of episode rewards
    """
    return run_experiment_with_config(cfg, agents, experiment_type, tags)


def run_experiment(
    cfg: DictConfig, agents: Dict[str, Agent], experiment_dir: Path
) -> Dict[str, Dict[int, float]]:
    """Run the actual experiment.

    This function handles the core experiment logic without any logging or error handling.

    Args:
        cfg: Configuration for the experiment
        agents: Dictionary mapping agent names to agent instances
        experiment_dir: Directory to save experiment outputs

    Returns:
        Dictionary mapping agent names to dictionaries of episode rewards
    """
    start_time = time.time()
    # Save configuration
    path = Path(experiment_dir)
    with open(path / "configuration.yaml", "w") as f:
        OmegaConf.save(cfg, f)

    # Create environment
    base_env = PoolFlipEnv(possible_agents=agents, configuration=cfg.env)

    # Save agent mapping
    with open(path / "agents.json", "w") as f:
        json.dump(base_env.agent_index_to_name, f)

    with open(path / "agents_repr.json", "w") as f:
        json.dump([agent.__repr__() for agent in agents.values()], f)

    # Wrap environment with experiment saver
    env = ExperimentSaverParallelWrapper(base_env, path / "data")

    # Run experiment
    if cfg.early_stopping.enabled:
        early_stopping = EarlyStopping(
            patience=cfg.early_stopping.patience,
            min_delta=cfg.early_stopping.min_delta,
            enabled=cfg.early_stopping.enabled,
            window_size=cfg.early_stopping.window_size,
        )
    else:
        early_stopping = None
    experiment_runner = ExperimentRunner(
        env, cfg.save_every_k_episodes, path / "checkpoints", early_stopping
    )
    results = experiment_runner.run(cfg.num_episodes, cfg.num_steps)

    end_time = time.time()
    with open(path / "completed.txt", "w") as f:
        f.write(f"Completed in {end_time - start_time} seconds")

    # Log artifacts
    mlflow.log_artifact(path / "completed.txt")
    mlflow.log_artifact(path / "configuration.yaml")
    mlflow.log_artifact(path / "agents.json")
    mlflow.log_artifact(path / "agents_repr.json")
    mlflow.log_artifacts(path / "data", artifact_path="data")
    for agent_name in agents:
        mlflow.log_artifact(
            path / "checkpoints" / f"{agent_name}_initial_policy.pth",
            artifact_path="checkpoints",
        )
        mlflow.log_artifact(
            path / "checkpoints" / f"{agent_name}_final_policy.pth",
            artifact_path="checkpoints",
        )

    # Log early stopping status and parameters
    if early_stopping:
        early_stopping_status = early_stopping.get_status()
        mlflow.log_params(
            {
                "early_stopping_enabled": early_stopping_status["enabled"],
                "early_stopping_patience": early_stopping_status["patience"],
                "early_stopping_min_delta": early_stopping_status["min_delta"],
                "early_stopping_window_size": early_stopping_status["window_size"],
            }
        )

        # Log scalar metrics
        mlflow.log_metrics(
            {
                "early_stopping_triggered_episode": early_stopping_status[
                    "episode_triggered"
                ],
            }
        )

    return results