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

.gitignore

@@ -523,3 +523,6 @@ Network Trash Folder
 Temporary Items
 .apdisk
 *.pdf
+
+# plot directory
+poster_plots/

configs/centralized-final.yaml

@@ -23,7 +23,7 @@ morphology:
   morph_mode: CENTRALIZED
 
 experiment:
-  exp_name: "final-models/centralized/"
+  exp_name: "final-models-v2/centralized/"
   seed: 42
   torch_deterministic: true
   cuda: true
@@ -34,8 +34,8 @@ logging:
   save_checkpoints: true
   upload_final_model: true
   upload_checkpoints: true
-  checkpoint_frequency: 20
+  checkpoint_frequency: 10
-  wandb_project_name: "final-models"
+  wandb_project_name: "final-models-v2"
 
 evaluation:
   evaluate_checkpoints: true

configs/evaluation/poster.yaml

@@ -9,12 +9,14 @@ eval_seed: 0
 # Cross-model comparison settings
 # We use 10 episodes to get a more robust average for the final poster results.
 comparison_base_seed: 0
-comparison_num_episodes: 2
+comparison_num_episodes: 10
 comparison_output_csv: "runs/evaluation/comparison.csv"
 
 # Paths to the .cleanrl_model files to be compared (relative to workspace root).
 comparison_models:
-  - "runs/input-space-2-arms/2026-05-02/08-14-58/final_model.flax"
+  - "runs/final-v2-centralized/artifacts/12-19-01_checkpoint_v22/checkpoint_step_230.flax"
+  - "runs/final-v2-fully-conn/artifacts/14-02-00_checkpoint_v17/checkpoint_step_180.flax"
+  - "runs/final-v2-ring/artifacts/15-27-03_checkpoint_v21/checkpoint_step_220.flax"
 
 # Path to the morphologies to evaluate against.
 comparison_morphologies:

configs/fully-connected-final.yaml

@@ -26,7 +26,7 @@ morphology:
   morph_mode: FULLY_CONNECTED
 
 experiment:
-  exp_name: "final-models/fully-connected/"
+  exp_name: "final-models-v2/fully-connected/"
   seed: 42
   torch_deterministic: true
   cuda: true
@@ -37,8 +37,8 @@ logging:
   save_checkpoints: true
   upload_final_model: true
   upload_checkpoints: true
-  checkpoint_frequency: 20
+  checkpoint_frequency: 10
-  wandb_project_name: "final-models"
+  wandb_project_name: "final-models-v2"
 
 evaluation:
   evaluate_checkpoints: true

configs/ring-final.yaml

@@ -26,7 +26,7 @@ morphology:
   morph_mode: RING
 
 experiment:
-  exp_name: "final-models/ring/"
+  exp_name: "final-models-v2/ring/"
   seed: 42
   torch_deterministic: true
   cuda: true
@@ -37,8 +37,8 @@ logging:
   save_checkpoints: true
   upload_final_model: true
   upload_checkpoints: true
-  checkpoint_frequency: 20
+  checkpoint_frequency: 10
-  wandb_project_name: "final-models"
+  wandb_project_name: "final-models-v2"
 
 evaluation:
   evaluate_checkpoints: true

scripts/plots/analyze_convergence.py

@@ -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(
+                "Reward_Convergence_Timestep": reward_timestep,
-                    arch_data[Columns.REWARD], arch_data[Columns.TIMESTEPS]
+                "Reward_Convergence_Checkpoint_Idx": reward_checkpoint_idx,
-                ),
+                "Reward_Convergence_Checkpoint": reward_checkpoint,
-                "Velocity_Convergence_Timestep": _convergence_timestep(
+                "Velocity_Convergence_Timestep": velocity_timestep,
-                    arch_data[Columns.VELOCITY], arch_data[Columns.TIMESTEPS]
+                "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.")