385 lines
13 KiB
Python
385 lines
13 KiB
Python
"""
|
|
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
|
|
|
|
from plot_config import COLORS, apply_style, create_common_parser, LEGEND_KWARGS
|
|
|
|
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."""
|
|
|
|
CHECKPOINT = "checkpoint"
|
|
ARCH = "architecture"
|
|
TIMESTEPS = "trained_timesteps"
|
|
REWARD = "eval_return"
|
|
VELOCITY = "velocity"
|
|
EVAL_STEPS = "eval_steps"
|
|
FINAL_XY_DIST = "final_xy_dist"
|
|
INITIAL_XY_DIST = "initial_xy_dist"
|
|
REACHED_TARGET = "reached_target"
|
|
|
|
|
|
# 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/final-v2-centralized/checkpoint_evaluation.csv",
|
|
"decentralized fully connected": "runs/final-v2-fully-conn/checkpoint_evaluation.csv",
|
|
"decentralized ring-level": "runs/final-v2-ring/checkpoint_evaluation.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.CHECKPOINT,
|
|
Columns.TIMESTEPS,
|
|
Columns.REWARD,
|
|
Columns.INITIAL_XY_DIST,
|
|
Columns.FINAL_XY_DIST,
|
|
Columns.EVAL_STEPS,
|
|
]
|
|
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.VELOCITY] = (df[Columns.INITIAL_XY_DIST] - df[Columns.FINAL_XY_DIST]) / df[
|
|
Columns.EVAL_STEPS
|
|
]
|
|
df[Columns.ARCH] = arch_name
|
|
df[Columns.VELOCITY] = (df[Columns.INITIAL_XY_DIST] - df[Columns.FINAL_XY_DIST]) / df[
|
|
Columns.EVAL_STEPS
|
|
]
|
|
|
|
dfs.append(df)
|
|
|
|
return pd.concat(dfs, ignore_index=True) if dfs else pd.DataFrame()
|
|
|
|
|
|
def _convergence_timestep(
|
|
series: pd.Series, timesteps: pd.Series, checkpoints: pd.Series
|
|
) -> tuple[float, int, int]:
|
|
"""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
|
|
|
|
mask = smoothed >= threshold
|
|
first_idx = mask.idxmax()
|
|
|
|
return timesteps.loc[first_idx], first_idx, checkpoints.loc[first_idx]
|
|
|
|
|
|
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 = []
|
|
|
|
centralized_base = 0
|
|
|
|
for arch in df[Columns.ARCH].unique():
|
|
arch_data = df[df[Columns.ARCH] == arch].sort_values(Columns.TIMESTEPS)
|
|
|
|
reward_timestep, reward_checkpoint_idx, reward_checkpoint = _convergence_timestep(
|
|
arch_data[Columns.REWARD],
|
|
arch_data[Columns.TIMESTEPS],
|
|
arch_data[Columns.CHECKPOINT],
|
|
)
|
|
|
|
velocity_timestep, velocity_checkpoint_idx, velocity_checkpoint = _convergence_timestep(
|
|
arch_data[Columns.VELOCITY],
|
|
arch_data[Columns.TIMESTEPS],
|
|
arch_data[Columns.CHECKPOINT],
|
|
)
|
|
|
|
results.append(
|
|
{
|
|
"Architecture": arch,
|
|
"Reward_Convergence_Timestep": reward_timestep,
|
|
"Reward_Convergence_Checkpoint_Idx": reward_checkpoint_idx,
|
|
"Reward_Convergence_Checkpoint": reward_checkpoint,
|
|
"Velocity_Convergence_Timestep": velocity_timestep,
|
|
"Velocity_Convergence_Checkpoint_Idx": velocity_checkpoint_idx,
|
|
"Velocity_Convergence_Checkpoint": velocity_checkpoint,
|
|
}
|
|
)
|
|
|
|
if arch == "centralized 5 arms":
|
|
centralized_base = reward_checkpoint
|
|
else:
|
|
print(arch, "speedup:", 1 - reward_checkpoint / centralized_base)
|
|
|
|
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, output_dir: str, **kwargs
|
|
):
|
|
"""
|
|
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()
|
|
|
|
fig, ax = plt.subplots(figsize=kwargs.get("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 ""
|
|
if kwargs.get("show_titles", True):
|
|
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(**LEGEND_KWARGS, ncol=2)
|
|
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()
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
base_path = os.path.join(output_dir, os.path.splitext(output_filename)[0])
|
|
plt.savefig(f"{base_path}.png", dpi=300, bbox_inches="tight")
|
|
plt.savefig(f"{base_path}.svg", format="svg", bbox_inches="tight")
|
|
plt.close()
|
|
|
|
|
|
def plot_metric_curves(
|
|
df: pd.DataFrame, metric_col: str, title: str, output_filename: str, output_dir: str, **kwargs
|
|
):
|
|
"""
|
|
Saves a line plot of the given metric over training timesteps for every architecture.
|
|
"""
|
|
fig, ax = plt.subplots(figsize=kwargs.get("figsize", (12, 7)))
|
|
|
|
for arch in df[Columns.ARCH].unique():
|
|
arch_data = df[df[Columns.ARCH] == arch].sort_values(Columns.TIMESTEPS)
|
|
color_key = arch.split()[0].upper() if isinstance(arch, str) else "UNKNOWN"
|
|
color = COLORS.get(color_key, "#888888")
|
|
ax.plot(
|
|
arch_data[Columns.TIMESTEPS],
|
|
arch_data[metric_col],
|
|
label=arch,
|
|
marker="o",
|
|
markersize=4,
|
|
alpha=0.8,
|
|
color=color,
|
|
)
|
|
|
|
title_suffix = " (DUMMY DATA)" if USING_DUMMY_DATA else ""
|
|
if kwargs.get("show_titles", True):
|
|
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(**LEGEND_KWARGS, ncol=len(df[Columns.ARCH].unique()))
|
|
ax.grid(True, linestyle="--", alpha=0.6)
|
|
ax.set_xlim(left=0)
|
|
ax.set_ylim(bottom=0)
|
|
|
|
plt.tight_layout()
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
base_path = os.path.join(output_dir, os.path.splitext(output_filename)[0])
|
|
plt.savefig(f"{base_path}.png", dpi=300, bbox_inches="tight")
|
|
plt.savefig(f"{base_path}.svg", format="svg", bbox_inches="tight")
|
|
plt.close()
|
|
|
|
|
|
def plot_results(df: pd.DataFrame, results: pd.DataFrame, output_dir: str, **kwargs):
|
|
"""Generates and saves all analysis plots."""
|
|
plot_grouped_convergence_chart(
|
|
results, output_filename="convergence_comparison.png", output_dir=output_dir, **kwargs
|
|
)
|
|
plot_metric_curves(
|
|
df,
|
|
Columns.REWARD,
|
|
"Training Progress: Accumulated Reward",
|
|
"progress_reward_curves.png",
|
|
output_dir=output_dir,
|
|
**kwargs,
|
|
)
|
|
plot_metric_curves(
|
|
df,
|
|
Columns.VELOCITY,
|
|
"Training Progress: Velocity",
|
|
"progress_velocity_curves.png",
|
|
output_dir=output_dir,
|
|
**kwargs,
|
|
)
|
|
|
|
|
|
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(output_dir: str, **kwargs):
|
|
"""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)
|
|
print(
|
|
results[
|
|
["Architecture", "Reward_Convergence_Checkpoint_Idx", "Reward_Convergence_Checkpoint"]
|
|
]
|
|
)
|
|
|
|
plot_results(df, results, output_dir, **kwargs)
|
|
logger.info("Analysis complete. Plots saved to disk.")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = create_common_parser(description="Analyze training convergence.")
|
|
args = parser.parse_args()
|
|
|
|
apply_style(font_size=args.font_size)
|
|
run_analysis(
|
|
output_dir=args.output_dir,
|
|
show_titles=args.show_titles,
|
|
figsize=(args.fig_width, args.fig_height),
|
|
)
|