1
Fork 0

feat(convergence analysis): updated convergence analysis script to also print out index of found item

This commit is contained in:
Robin Meersman 2026-05-15 20:52:06 +02:00
parent fc80799899
commit 7334565d69
6 changed files with 60 additions and 19 deletions

View file

@ -44,6 +44,7 @@ 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 = "total_trained_timesteps"
REWARD = "accumulated_reward"
@ -108,7 +109,18 @@ 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.
"""
<<<<<<< Updated upstream
required = [Columns.TIMESTEPS, Columns.REWARD, Columns.VELOCITY]
=======
required = [
Columns.CHECKPOINT,
Columns.TIMESTEPS,
Columns.REWARD,
Columns.INITIAL_XY_DIST,
Columns.FINAL_XY_DIST,
Columns.EVAL_STEPS,
]
>>>>>>> Stashed changes
dfs = []
for arch_name, filepath in file_mapping.items():
@ -130,11 +142,17 @@ def load_metrics(file_mapping: dict[str, str]) -> pd.DataFrame:
return pd.concat(dfs, ignore_index=True) if dfs else pd.DataFrame()
def _convergence_timestep(series: pd.Series, timesteps: pd.Series) -> float:
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
return timesteps[smoothed >= threshold].iloc[0]
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:
@ -147,15 +165,27 @@ def analyze_convergence(df: pd.DataFrame) -> pd.DataFrame:
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": _convergence_timestep(
arch_data[Columns.REWARD], arch_data[Columns.TIMESTEPS]
),
"Velocity_Convergence_Timestep": _convergence_timestep(
arch_data[Columns.VELOCITY], arch_data[Columns.TIMESTEPS]
),
"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,
}
)
@ -319,6 +349,12 @@ def run_analysis(output_dir: str, **kwargs):
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.")