| """Utilities for saving training/testing visualizations."""
|
|
|
| from __future__ import annotations
|
|
|
| from typing import Iterable, Optional
|
|
|
| import numpy as np
|
|
|
| import matplotlib
|
|
|
| matplotlib.use("Agg")
|
| import matplotlib.pyplot as plt
|
|
|
|
|
| def _to_numpy(values: Iterable) -> np.ndarray:
|
| if values is None:
|
| return np.array([])
|
| if isinstance(values, np.ndarray):
|
| return values
|
| try:
|
| import torch
|
|
|
| if torch.is_tensor(values):
|
| return values.detach().cpu().numpy()
|
| except ImportError:
|
| pass
|
| return np.asarray(values)
|
|
|
|
|
| def plot_flow_time_loss(
|
| time_loss_pairs: Iterable[tuple[float, float]],
|
| save_path: str,
|
| *,
|
| num_bins: int = 10,
|
| ) -> None:
|
| pairs = list(time_loss_pairs)
|
| if not pairs:
|
| return
|
| times = _to_numpy([pair[0] for pair in pairs]).astype(float)
|
| losses = _to_numpy([pair[1] for pair in pairs]).astype(float)
|
|
|
| plt.figure(figsize=(6, 4))
|
| plt.scatter(times, losses, alpha=0.3, s=12, label="samples")
|
|
|
| bins = np.linspace(0.0, 1.0, num_bins + 1)
|
| bin_indices = np.digitize(times, bins) - 1
|
| bin_centers = 0.5 * (bins[:-1] + bins[1:])
|
| bin_means = []
|
| for idx in range(num_bins):
|
| mask = bin_indices == idx
|
| if np.any(mask):
|
| bin_means.append(np.mean(losses[mask]))
|
| else:
|
| bin_means.append(np.nan)
|
| plt.plot(bin_centers, bin_means, color="red", linewidth=2, label="binned mean")
|
|
|
| plt.xlim(0.0, 1.0)
|
| plt.xlabel("t")
|
| plt.ylabel("loss")
|
| plt.title("Flow time loss vs time")
|
| plt.legend()
|
| plt.tight_layout()
|
| plt.savefig(save_path)
|
| plt.close()
|
|
|
|
|
| def plot_uncertainty_calibration(
|
| predicted: Iterable[float],
|
| empirical: Iterable[float],
|
| save_path: str,
|
| *,
|
| num_bins: int = 10,
|
| ) -> None:
|
| predicted_arr = _to_numpy(predicted).astype(float)
|
| empirical_arr = _to_numpy(empirical).astype(float)
|
| if predicted_arr.size == 0 or empirical_arr.size == 0:
|
| return
|
|
|
| min_val = np.min(predicted_arr)
|
| max_val = np.max(predicted_arr)
|
| if np.isclose(min_val, max_val):
|
| max_val = min_val + 1e-6
|
|
|
| bins = np.linspace(min_val, max_val, num_bins + 1)
|
| bin_indices = np.digitize(predicted_arr, bins) - 1
|
| bin_means_pred = []
|
| bin_means_emp = []
|
| for idx in range(num_bins):
|
| mask = bin_indices == idx
|
| if np.any(mask):
|
| bin_means_pred.append(np.mean(predicted_arr[mask]))
|
| bin_means_emp.append(np.mean(empirical_arr[mask]))
|
|
|
| plt.figure(figsize=(5, 5))
|
| plt.plot(bin_means_pred, bin_means_emp, marker="o", label="calibration")
|
| line_min = min(np.min(predicted_arr), np.min(empirical_arr))
|
| line_max = max(np.max(predicted_arr), np.max(empirical_arr))
|
| plt.plot([line_min, line_max], [line_min, line_max], linestyle="--", color="gray", label="y=x")
|
| plt.xlabel("Mean predicted uncertainty")
|
| plt.ylabel("Mean empirical error")
|
| plt.title("Uncertainty calibration")
|
| plt.legend()
|
| plt.tight_layout()
|
| plt.savefig(save_path)
|
| plt.close()
|
|
|
|
|
| def plot_training_trends(
|
| steps: Iterable[int],
|
| loss_fm: Iterable[float],
|
| loss_u: Optional[Iterable[float]],
|
| mean_pred_uncertainty: Iterable[float],
|
| mean_residual_magnitude: Iterable[float],
|
| save_path: str,
|
| ) -> None:
|
| steps_arr = _to_numpy(steps)
|
| plt.figure(figsize=(7, 4))
|
| plt.plot(steps_arr, _to_numpy(loss_fm), label="loss_fm")
|
| if loss_u is not None:
|
| loss_u_arr = _to_numpy(loss_u)
|
| if loss_u_arr.size:
|
| plt.plot(steps_arr, loss_u_arr, label="loss_u")
|
| plt.plot(steps_arr, _to_numpy(mean_pred_uncertainty), label="mean predicted uncertainty")
|
| plt.plot(steps_arr, _to_numpy(mean_residual_magnitude), label="mean residual magnitude")
|
| plt.xlabel("global step")
|
| plt.ylabel("value")
|
| plt.title("Training trends")
|
| plt.legend()
|
| plt.tight_layout()
|
| plt.savefig(save_path)
|
| plt.close()
|
|
|
|
|
| def plot_uncertainty_heatmap(
|
| uncertainty: Iterable,
|
| save_path: str,
|
| *,
|
| cmap: str = "viridis",
|
| ) -> None:
|
| uncertainty_arr = _to_numpy(uncertainty).astype(float)
|
| if uncertainty_arr.size == 0:
|
| return
|
| plt.figure(figsize=(6, 4))
|
| plt.imshow(uncertainty_arr, aspect="auto", cmap=cmap)
|
| plt.colorbar(label="s_hat")
|
| plt.xlabel("Action dimension")
|
| plt.ylabel("Timestep")
|
| plt.title("Uncertainty heatmap")
|
| plt.tight_layout()
|
| plt.savefig(save_path)
|
| plt.close()
|
|
|
|
|
| def plot_refinement_effect(
|
| mean_s_hat: Iterable[float],
|
| mean_residual: Optional[Iterable[float]],
|
| save_path: str,
|
| ) -> None:
|
| mean_s_hat_arr = _to_numpy(mean_s_hat)
|
| if mean_s_hat_arr.size == 0:
|
| return
|
| iterations = np.arange(1, len(mean_s_hat_arr) + 1)
|
| plt.figure(figsize=(6, 4))
|
| plt.plot(iterations, mean_s_hat_arr, marker="o", label="mean s_hat")
|
| if mean_residual is not None:
|
| mean_residual_arr = _to_numpy(mean_residual)
|
| if mean_residual_arr.size:
|
| plt.plot(iterations, mean_residual_arr, marker="x", label="mean residual magnitude")
|
| plt.xlabel("Refinement iteration")
|
| plt.ylabel("value")
|
| plt.title("Refinement effect")
|
| plt.legend()
|
| plt.tight_layout()
|
| plt.savefig(save_path)
|
| plt.close()
|
|
|
|
|
| def plot_task_summary(
|
| task_labels: Iterable[str],
|
| success_rate: Iterable[float],
|
| mean_uncertainty: Iterable[float],
|
| mean_failure_uncertainty: Iterable[float],
|
| save_path: str,
|
| ) -> None:
|
| labels = list(task_labels)
|
| success_arr = _to_numpy(success_rate).astype(float)
|
| mean_uncertainty_arr = _to_numpy(mean_uncertainty).astype(float)
|
| mean_failure_arr = _to_numpy(mean_failure_uncertainty).astype(float)
|
|
|
| x = np.arange(len(labels))
|
| width = 0.25
|
| plt.figure(figsize=(10, 4))
|
| plt.bar(x - width, success_arr, width=width, label="success rate")
|
| plt.bar(x, mean_uncertainty_arr, width=width, label="mean uncertainty")
|
| plt.bar(x + width, mean_failure_arr, width=width, label="mean uncertainty (failure)")
|
| plt.xticks(x, labels, rotation=45, ha="right")
|
| plt.ylabel("value")
|
| plt.title("Per-task summary")
|
| plt.legend()
|
| plt.tight_layout()
|
| plt.savefig(save_path)
|
| plt.close()
|
|
|
