diff --git a/scripts/evaluate_models.py b/scripts/evaluate_models.py index 12c2d0e..a64e972 100644 --- a/scripts/evaluate_models.py +++ b/scripts/evaluate_models.py @@ -19,7 +19,6 @@ import yaml matplotlib.use("Agg") import matplotlib.pyplot as plt -from matplotlib import colors from agents.appo_agent import APPOAgent from agents.dcmappo_agent import DCMAPPOAgent @@ -39,15 +38,18 @@ from agents.td3_agent import TD3Agent from envs.edge_vsl_env import SUMOEdgeVSLEnvironment from envs.reward_system import REWARD_COMPONENT_COLUMNS, REWARD_COMPONENT_LABELS from utils.config import get_agent_config +from utils.heatmap_plotting import ( + build_action_panel, + build_occupancy_panel, + build_speed_panel, + save_heatmap_panels, +) from utils.run_dirs import find_shared_config_path, resolve_checkpoint_root MODEL_ORDER = ["ppo", "gpro", "appo", "mappo", "tcamappo", "dcmappo", "dqn", "madqn", "ddqn", "qmix", "dcqmix", "ddpg", "sac", "td3", "sctd3"] BASELINE_NAME = "no_control" EVAL_ORDER = [BASELINE_NAME] + MODEL_ORDER -HEATMAP_SPEED_RANGE_KMH = (40.0, 110.0) -HEATMAP_OCCUPANCY_RANGE = (0.0, 35.0) -HEATMAP_ACTION_LEVELS_KMH = [40.0, 60.0, 80.0, 100.0, 110.0] MODEL_LABELS = { BASELINE_NAME: "NO_CONTROL", "ppo": "PPO", @@ -917,19 +919,6 @@ def plot_summary_bars(summary_df: pd.DataFrame, output_dir: str): def plot_model_heatmaps(edge_df: pd.DataFrame, detector_df: pd.DataFrame, output_dir: str): heatmap_dir = os.path.join(output_dir, "heatmaps") os.makedirs(heatmap_dir, exist_ok=True) - speed_cmap = plt.get_cmap("RdYlGn").copy() - speed_cmap.set_bad(color="#d9d9d9") - occ_cmap = plt.get_cmap("magma").copy() - occ_cmap.set_bad(color="#d9d9d9") - action_cmap = plt.get_cmap("viridis", len(HEATMAP_ACTION_LEVELS_KMH)).copy() - action_cmap.set_bad(color="#d9d9d9") - action_boundaries = [HEATMAP_ACTION_LEVELS_KMH[0] - 10.0] - action_boundaries.extend( - (left + right) / 2.0 - for left, right in zip(HEATMAP_ACTION_LEVELS_KMH[:-1], HEATMAP_ACTION_LEVELS_KMH[1:]) - ) - action_boundaries.append(HEATMAP_ACTION_LEVELS_KMH[-1] + 10.0) - action_norm = colors.BoundaryNorm(action_boundaries, action_cmap.N, clip=True) for model_name in EVAL_ORDER: detector_model_df = detector_df[detector_df["model"] == model_name] @@ -951,70 +940,39 @@ def plot_model_heatmaps(edge_df: pd.DataFrame, detector_df: pd.DataFrame, output .sort_values("edge_index") ) ordered_edge_ids = edge_order["edge_id"].tolist() - action_grid = edge_model_df.pivot(index="edge_id", columns="step", values="action_speed_kmh").reindex(ordered_edge_ids).values - - fig, axes = plt.subplots(1, 3, figsize=(21, 7), gridspec_kw={"width_ratios": [1.3, 0.8, 1.3]}) - - speed_im = axes[0].imshow( - np.ma.masked_invalid(speed_grid), - aspect="auto", - origin="lower", - cmap=speed_cmap, - vmin=HEATMAP_SPEED_RANGE_KMH[0], - vmax=HEATMAP_SPEED_RANGE_KMH[1], - ) - axes[0].set_title(f"{MODEL_LABELS[model_name]} Measured Speed (km/h)") - axes[0].set_xlabel("Step") - axes[0].set_ylabel("Detector Cell (bottom=upstream, top=downstream)") - plt.colorbar( - speed_im, - ax=axes[0], - fraction=0.046, - pad=0.04, - ticks=HEATMAP_ACTION_LEVELS_KMH, + action_plot_df = edge_model_df.copy() + if "action_applied" in action_plot_df.columns: + action_plot_df.loc[~action_plot_df["action_applied"].astype(bool), "action_speed_kmh"] = np.nan + action_grid = ( + action_plot_df.pivot(index="edge_id", columns="step", values="action_speed_kmh") + .reindex(ordered_edge_ids) + .values ) - action_im = axes[1].imshow( - np.ma.masked_invalid(action_grid), - aspect="auto", - origin="lower", - cmap=action_cmap, - norm=action_norm, + panels = [ + build_speed_panel( + speed_grid, + ordered_cell_ids, + f"{MODEL_LABELS[model_name]} Measured Speed (km/h)", + "Detector Cell (bottom=upstream, top=downstream)", + ), + build_action_panel( + action_grid, + ordered_edge_ids, + f"{MODEL_LABELS[model_name]} Applied VSL (km/h)", + ), + build_occupancy_panel( + occ_grid, + ordered_cell_ids, + f"{MODEL_LABELS[model_name]} Occupancy (%)", + "Detector Cell (bottom=upstream, top=downstream)", + ), + ] + save_heatmap_panels( + os.path.join(heatmap_dir, f"{model_name}_heatmaps.png"), + panels, + xlabel="Decision Step", ) - axes[1].set_title(f"{MODEL_LABELS[model_name]} Applied VSL (km/h)") - axes[1].set_xlabel("Step") - axes[1].set_ylabel("Controlled Edge") - plt.colorbar( - action_im, - ax=axes[1], - fraction=0.046, - pad=0.04, - ticks=HEATMAP_ACTION_LEVELS_KMH, - boundaries=action_boundaries, - ) - - occ_im = axes[2].imshow( - np.ma.masked_invalid(occ_grid), - aspect="auto", - origin="lower", - cmap=occ_cmap, - vmin=HEATMAP_OCCUPANCY_RANGE[0], - vmax=HEATMAP_OCCUPANCY_RANGE[1], - ) - axes[2].set_title(f"{MODEL_LABELS[model_name]} Occupancy (%)") - axes[2].set_xlabel("Step") - axes[2].set_ylabel("Detector Cell (bottom=upstream, top=downstream)") - plt.colorbar( - occ_im, - ax=axes[2], - fraction=0.046, - pad=0.04, - ticks=np.arange(0.0, 36.0, 5.0), - ) - - plt.tight_layout() - plt.savefig(os.path.join(heatmap_dir, f"{model_name}_heatmaps.png"), dpi=160) - plt.close() def _format_metric(value: float, fmt: str) -> str: diff --git a/utils/episode_artifacts.py b/utils/episode_artifacts.py index 390b8c3..c1fc2ab 100644 --- a/utils/episode_artifacts.py +++ b/utils/episode_artifacts.py @@ -9,9 +9,14 @@ import matplotlib import numpy as np matplotlib.use("Agg") -import matplotlib.pyplot as plt from envs.reward_system import REWARD_COMPONENT_COLUMNS +from utils.heatmap_plotting import ( + build_action_panel, + build_density_panel, + build_speed_panel, + save_heatmap_panels, +) def _safe_float(value, default: float = float("nan")) -> float: @@ -50,7 +55,8 @@ def _normalize_step_rows(episode: int, episode_metrics: Sequence[Dict]) -> Tuple action_speeds = list(info.get("edge_speeds_kmh", [])) measured_speeds_ms = list(info.get("edge_speeds_ms", [])) occupancies = list(info.get("edge_occupancies", [])) - edge_count = max(len(action_speeds), len(measured_speeds_ms), len(occupancies)) + action_applied_mask = list(info.get("action_applied_mask", [])) + edge_count = max(len(action_speeds), len(measured_speeds_ms), len(occupancies), len(action_applied_mask)) if not edge_ids: edge_ids = [f"edge_{idx:02d}" for idx in range(edge_count)] @@ -63,6 +69,7 @@ def _normalize_step_rows(episode: int, episode_metrics: Sequence[Dict]) -> Tuple "edge_index": edge_idx, "edge_id": edge_ids[edge_idx], "action_speed_kmh": _safe_float(action_speeds[edge_idx] if edge_idx < len(action_speeds) else np.nan), + "action_applied": bool(action_applied_mask[edge_idx]) if edge_idx < len(action_applied_mask) else True, "measured_speed_kmh": _safe_float( measured_speeds_ms[edge_idx] * 3.6 if edge_idx < len(measured_speeds_ms) else np.nan ), @@ -213,64 +220,87 @@ def _build_detector_rows_from_xml(log_dir: str, episode: int) -> List[Dict]: return detector_rows -def _plot_episode_heatmap(path: str, detector_rows: Sequence[Dict], title_prefix: str): - if not detector_rows: +def _plot_episode_heatmap( + path: str, + edge_rows: Sequence[Dict], + edge_ids: Sequence[str], + detector_rows: Sequence[Dict], + title_prefix: str, +): + has_action = bool(edge_rows and edge_ids) + has_detector = bool(detector_rows) + if not has_action and not has_detector: return - step_values = sorted({int(row["step"]) for row in detector_rows}) - ordered_cells = [] - seen_cells = set() - for row in sorted( - detector_rows, - key=lambda item: (int(item["cell_order"]), str(item["cell_id"])), - ): - cell_id = str(row["cell_id"]) - if cell_id in seen_cells: - continue - seen_cells.add(cell_id) - ordered_cells.append(cell_id) + step_values = sorted( + { + int(row["step"]) + for row in list(edge_rows) + list(detector_rows) + if row.get("step") is not None + } + ) + if not step_values: + return - num_cells = len(ordered_cells) num_steps = len(step_values) step_to_col = {step: idx for idx, step in enumerate(step_values)} - cell_to_row = {cell_id: idx for idx, cell_id in enumerate(ordered_cells)} + panels = [] + if has_action: + ordered_edge_ids = list(edge_ids) + edge_to_row = {edge_id: idx for idx, edge_id in enumerate(ordered_edge_ids)} + action_grid = np.full((len(ordered_edge_ids), num_steps), np.nan, dtype=np.float32) + for row in edge_rows: + edge_id = str(row["edge_id"]) + if edge_id not in edge_to_row: + continue + if not bool(row.get("action_applied", True)): + continue + row_idx = edge_to_row[edge_id] + col_idx = step_to_col[int(row["step"])] + action_grid[row_idx, col_idx] = _safe_float(row["action_speed_kmh"]) + panels.append(build_action_panel(action_grid, ordered_edge_ids, f"{title_prefix} Applied VSL (km/h)")) - speed_grid = np.full((num_cells, num_steps), np.nan, dtype=np.float32) - density_grid = np.full((num_cells, num_steps), np.nan, dtype=np.float32) + if has_detector: + ordered_cells = [] + seen_cells = set() + for row in sorted( + detector_rows, + key=lambda item: (int(item["cell_order"]), str(item["cell_id"])), + ): + cell_id = str(row["cell_id"]) + if cell_id in seen_cells: + continue + seen_cells.add(cell_id) + ordered_cells.append(cell_id) - for row in detector_rows: - row_idx = cell_to_row[str(row["cell_id"])] - col_idx = step_to_col[int(row["step"])] - speed_grid[row_idx, col_idx] = _safe_float(row["measured_speed_kmh"]) - density_grid[row_idx, col_idx] = _safe_float(row["density_vehpkm"]) + num_cells = len(ordered_cells) + cell_to_row = {cell_id: idx for idx, cell_id in enumerate(ordered_cells)} + speed_grid = np.full((num_cells, num_steps), np.nan, dtype=np.float32) + density_grid = np.full((num_cells, num_steps), np.nan, dtype=np.float32) - fig, axes = plt.subplots(1, 2, figsize=(18, 7), sharex=True, sharey=True) - plots = [ - (speed_grid, "RdYlGn", "Measured Speed (km/h)"), - (density_grid, "YlOrRd", "Density (veh/km)"), - ] + for row in detector_rows: + row_idx = cell_to_row[str(row["cell_id"])] + col_idx = step_to_col[int(row["step"])] + speed_grid[row_idx, col_idx] = _safe_float(row["measured_speed_kmh"]) + density_grid[row_idx, col_idx] = _safe_float(row["density_vehpkm"]) - for ax, (grid, cmap, title) in zip(axes, plots): - image = ax.imshow( - np.ma.masked_invalid(grid), - aspect="auto", - origin="lower", - cmap=cmap, - interpolation="nearest", - resample=False, - ) - ax.set_title(f"{title_prefix} {title}") - ax.set_xlabel("Decision Step") - ax.set_ylabel("Detector Cell") - tick_step = max(num_cells // 12, 1) - tick_idx = np.arange(0, num_cells, tick_step) - ax.set_yticks(tick_idx) - ax.set_yticklabels([ordered_cells[idx] for idx in tick_idx], fontsize=8) - plt.colorbar(image, ax=ax, fraction=0.046, pad=0.04) + plots = [ + build_speed_panel( + speed_grid, + ordered_cells, + f"{title_prefix} Measured Speed (km/h)", + "Detector Cell", + ), + build_density_panel( + density_grid, + ordered_cells, + f"{title_prefix} Density (veh/km)", + "Detector Cell", + ), + ] + panels = plots[:1] + panels + plots[1:] - plt.tight_layout() - plt.savefig(path, dpi=160) - plt.close(fig) + save_heatmap_panels(path, panels, xlabel="Decision Step") def _write_summary(path: str, summary: Optional[Dict]): @@ -303,7 +333,7 @@ def _save_episode_bundle( detector_rows, fieldnames=list(detector_rows[0].keys()) if detector_rows else ["episode", "step", "cell_id"], ) - _plot_episode_heatmap(heatmap_path, detector_rows, title_prefix=f"Episode {episode}") + _plot_episode_heatmap(heatmap_path, edge_rows, edge_ids, detector_rows, title_prefix=f"Episode {episode}") _write_summary(summary_path, summary) diff --git a/utils/heatmap_plotting.py b/utils/heatmap_plotting.py new file mode 100644 index 0000000..0ffd238 --- /dev/null +++ b/utils/heatmap_plotting.py @@ -0,0 +1,149 @@ +"""Shared heatmap styling and rendering utilities.""" +from typing import Dict, Optional, Sequence + +import matplotlib +import numpy as np + +matplotlib.use("Agg") +import matplotlib.pyplot as plt +from matplotlib import colors + +HEATMAP_SPEED_RANGE_KMH = (0.0, 110.0) +HEATMAP_OCCUPANCY_RANGE = (0.0, 35.0) +HEATMAP_ACTION_LEVELS_KMH = [40.0, 60.0, 80.0, 100.0, 110.0] +HEATMAP_OCCUPANCY_TICKS = np.arange(0.0, 36.0, 5.0) +HEATMAP_BAD_COLOR = "#d9d9d9" + + +def _get_masked_cmap(name: str, levels: Optional[int] = None): + cmap = plt.get_cmap(name, levels).copy() if levels is not None else plt.get_cmap(name).copy() + cmap.set_bad(color=HEATMAP_BAD_COLOR) + return cmap + + +def get_action_boundaries(): + boundaries = [HEATMAP_ACTION_LEVELS_KMH[0] - 10.0] + boundaries.extend( + (left + right) / 2.0 + for left, right in zip(HEATMAP_ACTION_LEVELS_KMH[:-1], HEATMAP_ACTION_LEVELS_KMH[1:]) + ) + boundaries.append(HEATMAP_ACTION_LEVELS_KMH[-1] + 10.0) + return boundaries + + +def get_action_norm(): + action_cmap = _get_masked_cmap("viridis", len(HEATMAP_ACTION_LEVELS_KMH)) + return action_cmap, colors.BoundaryNorm(get_action_boundaries(), action_cmap.N, clip=True) + + +def build_action_panel(grid, row_labels: Sequence[str], title: str) -> Dict: + action_cmap, action_norm = get_action_norm() + return { + "grid": grid, + "row_labels": row_labels, + "title": title, + "ylabel": "Controlled Edge", + "cmap": action_cmap, + "norm": action_norm, + "width_ratio": 0.8, + "colorbar_kwargs": { + "ticks": HEATMAP_ACTION_LEVELS_KMH, + "boundaries": get_action_boundaries(), + }, + } + + +def build_speed_panel(grid, row_labels: Sequence[str], title: str, ylabel: str) -> Dict: + return { + "grid": grid, + "row_labels": row_labels, + "title": title, + "ylabel": ylabel, + "cmap": _get_masked_cmap("RdYlGn"), + "vmin": HEATMAP_SPEED_RANGE_KMH[0], + "vmax": HEATMAP_SPEED_RANGE_KMH[1], + "width_ratio": 1.3, + "colorbar_kwargs": { + "ticks": HEATMAP_ACTION_LEVELS_KMH, + }, + } + + +def build_occupancy_panel(grid, row_labels: Sequence[str], title: str, ylabel: str) -> Dict: + return { + "grid": grid, + "row_labels": row_labels, + "title": title, + "ylabel": ylabel, + "cmap": _get_masked_cmap("magma"), + "vmin": HEATMAP_OCCUPANCY_RANGE[0], + "vmax": HEATMAP_OCCUPANCY_RANGE[1], + "width_ratio": 1.3, + "colorbar_kwargs": { + "ticks": HEATMAP_OCCUPANCY_TICKS, + }, + } + + +def build_density_panel(grid, row_labels: Sequence[str], title: str, ylabel: str) -> Dict: + return { + "grid": grid, + "row_labels": row_labels, + "title": title, + "ylabel": ylabel, + "cmap": _get_masked_cmap("YlOrRd"), + "width_ratio": 1.3, + "colorbar_kwargs": {}, + } + + +def save_heatmap_panels( + path: str, + panels: Sequence[Optional[Dict]], + xlabel: str = "Decision Step", + dpi: int = 160, +): + valid_panels = [panel for panel in panels if panel and panel.get("grid") is not None] + if not valid_panels: + return + + width_ratios = [float(panel.get("width_ratio", 1.0)) for panel in valid_panels] + fig, axes = plt.subplots( + 1, + len(valid_panels), + figsize=(7 * len(valid_panels), 7), + gridspec_kw={"width_ratios": width_ratios}, + ) + if len(valid_panels) == 1: + axes = [axes] + + for ax, panel in zip(axes, valid_panels): + grid = np.asarray(panel["grid"], dtype=np.float32) + image = ax.imshow( + np.ma.masked_invalid(grid), + aspect="auto", + origin="lower", + cmap=panel["cmap"], + interpolation="nearest", + resample=False, + vmin=panel.get("vmin"), + vmax=panel.get("vmax"), + norm=panel.get("norm"), + ) + ax.set_title(str(panel["title"])) + ax.set_xlabel(xlabel) + ax.set_ylabel(str(panel["ylabel"])) + + row_labels = list(panel.get("row_labels", [])) + if row_labels: + tick_step = max(len(row_labels) // 12, 1) + tick_idx = np.arange(0, len(row_labels), tick_step) + ax.set_yticks(tick_idx) + ax.set_yticklabels([row_labels[idx] for idx in tick_idx], fontsize=8) + + colorbar_kwargs = dict(panel.get("colorbar_kwargs", {})) + plt.colorbar(image, ax=ax, fraction=0.046, pad=0.04, **colorbar_kwargs) + + plt.tight_layout() + plt.savefig(path, dpi=dpi) + plt.close(fig)