diff --git a/scripts/analysis/analyze_convergence.py b/scripts/analysis/analyze_convergence.py new file mode 100644 index 0000000..c58a34b --- /dev/null +++ b/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()