Merge pull request #61 from SELab-3-2026/feat/final-plots

plot: final plotting code used to create plots showed on poster

.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_comparisons.py

@@ -6,18 +6,17 @@ Rate, Distance Remaining).
 """
 
 import os
-import pandas as pd
-import numpy as np
-import matplotlib.pyplot as plt
 
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
 from plot_config import (
-    COLORS,
-    apply_style,
-    BEST_PERFORMER_MARKER,
-    BEST_PERFORMER_TEXT,
     BEST_PERFORMER_COLOR,
-    create_common_parser,
+    BEST_PERFORMER_TEXT,
+    COLORS,
     LEGEND_KWARGS,
+    apply_style,
+    create_common_parser,
 )
 
 
@@ -84,7 +83,8 @@ def plot_grouped_bar(
 
     fig, ax = plt.subplots(figsize=figsize)
     bar_width = 0.35
-    x_indices = np.arange(len(morphologies))
+    group_spacing = 1.3
+    x_indices = np.arange(len(morphologies)) * group_spacing
     all_bars = {}
     all_means = []
 
@@ -118,7 +118,7 @@ def plot_grouped_bar(
         )
         all_bars[arch] = (x_pos, means, stds, bars)
 
-    for m_idx, m in enumerate(morphologies):
+    for m_idx, _ in enumerate(morphologies):
         m_means = {arch: all_bars[arch][1][m_idx] for arch in architectures}
         best_arch = (
             max(m_means, key=m_means.get) if higher_is_better else min(m_means, key=m_means.get)
@@ -144,12 +144,13 @@ def plot_grouped_bar(
 
     x_ticks_pos = (
         x_indices
+        + bar_width  # center the label in the 3 bars
         + (bar_width / 2 if len(architectures) % 2 == 0 else 0)
         - (bar_width / 2 if len(architectures) == 2 else 0)
     )
     ax.set_xticks(x_ticks_pos)
     ax.set_xticklabels([f"{m} Arms" for m in morphologies])
-    ax.tick_params(axis="x", pad=25)  # More padding for the squares
+    ax.tick_params(axis="x")  # More padding for the squares
 
     # X-axis at zero
     ax.axhline(0, color="black", linewidth=1.5)
@@ -172,20 +173,7 @@ def plot_grouped_bar(
             plt.FuncFormatter(lambda x, _: f"{x:.2f}" if abs(x) < 10 else f"{x:.0f}")
         )
 
-    _add_square_placeholders(ax, x_ticks_pos, [f"{m} Arms" for m in morphologies])
+    ax.legend(**LEGEND_KWARGS, ncol=len(architectures))
-
-    # Add custom legend entry for best performer
-    ax.plot(
-        [],
-        [],
-        marker=BEST_PERFORMER_MARKER,
-        color="w",
-        markerfacecolor=BEST_PERFORMER_COLOR,
-        markersize=15,
-        # label="Best Performance",
-        ls="",
-    )
-    ax.legend(**LEGEND_KWARGS, ncol=len(architectures) + 1)
     ax.set_facecolor("white")
     fig.patch.set_facecolor("white")
 
@@ -305,22 +293,7 @@ def plot_grouped_bar_alt(
             plt.FuncFormatter(lambda x, _: f"{x:.2f}" if abs(x) < 10 else f"{x:.0f}")
         )
 
-    # In this alt plot, placeholders might be per architecture
+    ax.legend(**LEGEND_KWARGS, ncol=len(morphologies))
-    _add_square_placeholders(
-        ax, x_indices, [arch.replace("_", "\n").title() for arch in architectures]
-    )
-
-    ax.plot(
-        [],
-        [],
-        marker=BEST_PERFORMER_MARKER,
-        color="w",
-        markerfacecolor=BEST_PERFORMER_COLOR,
-        markersize=15,
-        # label="Best Performance",
-        ls="",
-    )
-    ax.legend(**LEGEND_KWARGS, ncol=len(morphologies) + 1)
     ax.set_facecolor("white")
     fig.patch.set_facecolor("white")
 
@@ -358,8 +331,8 @@ if __name__ == "__main__":
         plot_grouped_bar(
             df=df,
             metric_col="approx_max_velocity",
-            ylabel="Max Forward Velocity (cm/s)",
+            ylabel="",
-            title="Graceful Degradation: Velocity Across Morphologies",
+            title="Maximal forward velocity (in cm/s)",
             output_filename="poster_plot_velocity.png",
             output_dir=OUTPUT_DIR,
             higher_is_better=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 = "trained_timesteps"
     REWARD = "eval_return"
@@ -112,6 +113,7 @@ def load_metrics(file_mapping: dict[str, str]) -> pd.DataFrame:
     and returns the combined DataFrame with only the required columns.
     """
     required = [
+        Columns.CHECKPOINT,
         Columns.TIMESTEPS,
         Columns.REWARD,
         Columns.INITIAL_XY_DIST,
@@ -133,6 +135,9 @@ def load_metrics(file_mapping: dict[str, str]) -> pd.DataFrame:
             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
@@ -143,11 +148,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:
@@ -157,21 +168,40 @@ def analyze_convergence(df: pd.DataFrame) -> pd.DataFrame:
     """
     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": _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,
             }
         )
 
+        if arch == "centralized 5 arms":
+            centralized_base = reward_checkpoint
+        else:
+            print(arch, "speedup:", 1 - reward_checkpoint / centralized_base)
+
     return pd.DataFrame(results)
 
 
@@ -333,6 +363,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.")
 

scripts/plots/plot_config.py

@@ -6,7 +6,7 @@ import matplotlib.pyplot as plt
 COLORS = {
     "CENTRALIZED": "#0D567C",  # Blue
     "FULLY_CONNECTED": "#8C0E0F",  # Reddish
-    "RING_LEVEL": "#E1BA6D",  # Pale Yellow
+    "RING": "#FCB305",  # Pale Yellow
 }
 
 
@@ -42,7 +42,7 @@ BEST_PERFORMER_COLOR = "#D4AF37"  # Gold
 # Centralized Legend Configuration
 LEGEND_KWARGS = {
     "loc": "upper center",
-    "bbox_to_anchor": (0.5, -0.5),
+    "bbox_to_anchor": (0.5, -0.12),
     "frameon": False,
 }