feat: convergence plots

scripts/analysis/analyze_convergence.py

@@ -0,0 +1,296 @@
+"""
+Convergence Analysis Script for Poster Visualizations
+
+This script analyzes evaluation metrics from multiple training runs to determine
+the convergence point of different reinforcement learning architectures.
+
+Workflow:
+1. Loads evaluation data from the CSV files defined in FILE_MAPPING.
+2. Calculates a rolling average of the reward and velocity to smooth noise.
+3. Determines the convergence timestep for each metric (first time 95% of peak is reached).
+4. Generates a grouped bar chart comparing convergence speed and line plots of the raw curves.
+
+Usage:
+    uv run python scripts/analysis/analyze_convergence.py
+
+Note: For these metrics to be valid, the evaluation CSVs must be generated with
+exploration noise strictly disabled (e.g., taking the mean of the action distribution).
+"""
+
+import logging
+import os
+
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+from enum import Enum
+
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
+    datefmt="%Y-%m-%d %H:%M:%S",
+)
+logger = logging.getLogger(__name__)
+
+# --- Globals & Configuration ---
+USING_DUMMY_DATA = False
+SMOOTHING_WINDOW = 3
+CONVERGENCE_THRESHOLD = 0.95
+
+
+class Columns(str, Enum):
+    # ... (rest of the file remains same, just need to update plotting functions and obtain_data)
+    """Column names expected in every evaluation CSV."""
+
+    ARCH = "architecture"
+    TIMESTEPS = "total_trained_timesteps"
+    REWARD = "accumulated_reward"
+    VELOCITY = "velocity"
+
+
+# Maps architecture display names to the path of their evaluation CSV.
+# Update these paths once real evaluation data is available.
+FILE_MAPPING: dict[str, str] = {
+    "centralized 2 arms": "runs/dummy/dummy_centralized_2_arms.csv",
+    "centralized 5 arms": "runs/dummy/dummy_centralized_5_arms.csv",
+    "decentralized fully connected": "runs/dummy/dummy_decentralized_fully_connected.csv",
+    "decentralized ring-level": "runs/dummy/dummy_decentralized_ring-level.csv",
+    "decentralized segment-level": "runs/dummy/dummy_decentralized_segment-level.csv",
+}
+
+# Architecture profiles for dummy data generation: (max_reward, max_velocity, sigmoid_speed)
+_DUMMY_PROFILES: dict[str, tuple[float, float, float]] = {
+    "centralized 2 arms": (300, 0.8, 1.2),
+    "centralized 5 arms": (450, 1.1, 1.0),
+    "decentralized fully connected": (500, 1.3, 0.7),
+    "decentralized ring-level": (480, 1.2, 0.8),
+    "decentralized segment-level": (520, 1.4, 0.6),
+}
+
+
+def generate_dummy_csvs(file_mapping: dict[str, str]):
+    """
+    Generates one dummy CSV per architecture in FILE_MAPPING at their expected locations.
+    Skips any architecture without a defined profile.
+    """
+    checkpoints = list(range(100, 1100, 100))
+    timesteps = [cp * 10_000 for cp in checkpoints]
+
+    for arch, path in file_mapping.items():
+        if arch not in _DUMMY_PROFILES:
+            logger.warning(f"No dummy profile for '{arch}'. Skipping.")
+            continue
+
+        m_reward, m_vel, speed = _DUMMY_PROFILES[arch]
+
+        rows = []
+        for i, ts in enumerate(timesteps):
+            progress = 1 / (1 + np.exp(-speed * (i - 4)))
+            rows.append(
+                {
+                    Columns.TIMESTEPS: ts,
+                    Columns.REWARD: m_reward * progress + np.random.normal(0, 5),
+                    Columns.VELOCITY: m_vel * progress + np.random.normal(0, 0.02),
+                }
+            )
+
+        # Create parent directories if they don't exist
+        os.makedirs(os.path.dirname(path), exist_ok=True)
+
+        pd.DataFrame(rows).to_csv(path, index=False)
+        logger.info(f"Generated dummy CSV at expected path: {path}")
+
+
+def load_metrics(file_mapping: dict[str, str]) -> pd.DataFrame:
+    """
+    Loads one CSV per architecture, injects the architecture name as a column,
+    and returns the combined DataFrame with only the required columns.
+    """
+    required = [Columns.TIMESTEPS, Columns.REWARD, Columns.VELOCITY]
+    dfs = []
+
+    for arch_name, filepath in file_mapping.items():
+        if not os.path.exists(filepath):
+            logger.warning(f"File not found: '{filepath}'. Skipping.")
+            continue
+
+        df = pd.read_csv(filepath)
+
+        missing = [c for c in required if c not in df.columns]
+        if missing:
+            logger.warning(f"Missing columns {missing} in '{filepath}'. Skipping.")
+            continue
+
+        df = df[required].copy()
+        df[Columns.ARCH] = arch_name
+        dfs.append(df)
+
+    return pd.concat(dfs, ignore_index=True) if dfs else pd.DataFrame()
+
+
+def _convergence_timestep(series: pd.Series, timesteps: pd.Series) -> float:
+    """Returns the first timestep where the smoothed series reaches 95% of its peak."""
+    smoothed = series.rolling(window=SMOOTHING_WINDOW, min_periods=1).mean()
+    threshold = smoothed.max() * CONVERGENCE_THRESHOLD
+    return timesteps[smoothed >= threshold].iloc[0]
+
+
+def analyze_convergence(df: pd.DataFrame) -> pd.DataFrame:
+    """
+    For each architecture, determines the convergence timestep based on both
+    reward and velocity, returning one summary row per architecture.
+    """
+    results = []
+
+    for arch in df[Columns.ARCH].unique():
+        arch_data = df[df[Columns.ARCH] == arch].sort_values(Columns.TIMESTEPS)
+
+        results.append(
+            {
+                "Architecture": arch,
+                "Reward_Convergence_Timestep": _convergence_timestep(
+                    arch_data[Columns.REWARD], arch_data[Columns.TIMESTEPS]
+                ),
+                "Velocity_Convergence_Timestep": _convergence_timestep(
+                    arch_data[Columns.VELOCITY], arch_data[Columns.TIMESTEPS]
+                ),
+            }
+        )
+
+    return pd.DataFrame(results)
+
+
+def _add_bar_labels(bars, max_val: float):
+    """Annotates each bar with its value in white bold text, positioned inside."""
+    for bar in bars:
+        width = bar.get_width()
+        label = f"{width / 1e6:.1f}M" if width >= 1e6 else f"{width:,.0f}"
+        plt.text(
+            width - (max_val * 0.02),
+            bar.get_y() + bar.get_height() / 2,
+            label,
+            ha="right",
+            va="center",
+            fontsize=11,
+            color="white",
+            fontweight="bold",
+        )
+
+
+def plot_grouped_convergence_chart(results_df: pd.DataFrame, output_filename: str):
+    """
+    Saves a grouped horizontal bar chart comparing Reward and Velocity convergence timesteps
+    across all architectures.
+    """
+    sorted_df = results_df.sort_values("Reward_Convergence_Timestep", ascending=True)
+    architectures = sorted_df["Architecture"].tolist()
+    y_pos = np.arange(len(architectures))
+    bar_height = 0.35
+    max_val = sorted_df[
+        ["Reward_Convergence_Timestep", "Velocity_Convergence_Timestep"]
+    ].values.max()
+
+    _, ax = plt.subplots(figsize=(12, 8))
+
+    bars_reward = ax.barh(
+        y_pos + bar_height / 2,
+        sorted_df["Reward_Convergence_Timestep"],
+        height=bar_height,
+        label="Reward Convergence",
+        color="#1f77b4",
+    )
+    bars_velocity = ax.barh(
+        y_pos - bar_height / 2,
+        sorted_df["Velocity_Convergence_Timestep"],
+        height=bar_height,
+        label="Velocity Convergence",
+        color="#ff7f0e",
+    )
+
+    title_suffix = " (DUMMY DATA)" if USING_DUMMY_DATA else ""
+    ax.set_title(f"Comparison of Training Convergence Timesteps{title_suffix}", fontsize=20, pad=20)
+    ax.set_xlabel("Timesteps to Convergence (95% of peak)", fontsize=16)
+    ax.set_ylabel("Architecture", fontsize=16)
+    ax.set_yticks(y_pos)
+    ax.set_yticklabels(architectures, fontsize=14)
+    ax.tick_params(axis="x", labelsize=14)
+    ax.legend(fontsize=12, loc="lower right")
+    ax.set_xlim(left=0)
+    ax.spines["top"].set_visible(False)
+    ax.spines["right"].set_visible(False)
+
+    _add_bar_labels(bars_reward, max_val)
+    _add_bar_labels(bars_velocity, max_val)
+
+    plt.tight_layout()
+    plt.savefig(output_filename, format="png", dpi=300, bbox_inches="tight")
+    plt.close()
+
+
+def plot_metric_curves(df: pd.DataFrame, metric_col: str, title: str, output_filename: str):
+    """
+    Saves a line plot of the given metric over training timesteps for every architecture.
+    """
+    _, ax = plt.subplots(figsize=(12, 7))
+
+    for arch in df[Columns.ARCH].unique():
+        arch_data = df[df[Columns.ARCH] == arch].sort_values(Columns.TIMESTEPS)
+        ax.plot(
+            arch_data[Columns.TIMESTEPS],
+            arch_data[metric_col],
+            label=arch,
+            marker="o",
+            markersize=4,
+            alpha=0.8,
+        )
+
+    title_suffix = " (DUMMY DATA)" if USING_DUMMY_DATA else ""
+    ax.set_title(f"{title}{title_suffix}", fontsize=18, pad=20)
+    ax.set_xlabel("Training Timesteps", fontsize=14)
+    ax.set_ylabel(metric_col.replace("_", " ").title(), fontsize=14)
+    ax.legend(bbox_to_anchor=(1.05, 1), loc="upper left", fontsize=10)
+    ax.grid(True, linestyle="--", alpha=0.6)
+    ax.set_xlim(left=0)
+    ax.set_ylim(bottom=0)
+
+    plt.tight_layout()
+    plt.savefig(output_filename, format="png", dpi=300, bbox_inches="tight")
+    plt.close()
+
+
+def plot_results(df: pd.DataFrame, results: pd.DataFrame):
+    """Generates and saves all analysis plots."""
+    plot_grouped_convergence_chart(results, output_filename="convergence_comparison.png")
+    plot_metric_curves(
+        df, Columns.REWARD, "Training Progress: Accumulated Reward", "progress_reward_curves.png"
+    )
+    plot_metric_curves(
+        df, Columns.VELOCITY, "Training Progress: Velocity", "progress_velocity_curves.png"
+    )
+
+
+def obtain_data() -> pd.DataFrame:
+    """Resolves the file mapping, falling back to generated dummy CSVs if needed."""
+    global USING_DUMMY_DATA
+    if not any(os.path.exists(p) for p in FILE_MAPPING.values()):
+        logger.info("No real evaluation files found. Generating dummy CSVs at expected locations.")
+        generate_dummy_csvs(FILE_MAPPING)
+        USING_DUMMY_DATA = True
+
+    return load_metrics(FILE_MAPPING)
+
+
+def run_analysis():
+    """Orchestrates data loading, convergence analysis, and plot generation."""
+    df = obtain_data()
+    if df.empty:
+        logger.error("No data found to analyze.")
+        return
+
+    results = analyze_convergence(df)
+    plot_results(df, results)
+    logger.info("Analysis complete. Plots saved to disk.")
+
+
+if __name__ == "__main__":
+    run_analysis()