| from __future__ import annotations | |
| import argparse | |
| import csv | |
| import math | |
| import os | |
| import statistics | |
| import sys | |
| from collections import defaultdict | |
| from pathlib import Path | |
| SCRIPT_ROOT = Path(__file__).resolve().parents[1] | |
| if str(SCRIPT_ROOT) not in sys.path: | |
| sys.path.insert(0, str(SCRIPT_ROOT)) | |
| from common import FINAL_ROOT | |
| MPLCONFIGDIR = FINAL_ROOT / ".cache" / "matplotlib" | |
| MPLCONFIGDIR.mkdir(parents=True, exist_ok=True) | |
| os.environ.setdefault("MPLCONFIGDIR", str(MPLCONFIGDIR)) | |
| import matplotlib | |
| matplotlib.use("Agg") | |
| import matplotlib.pyplot as plt | |
| ALGO_ORDER = ["hamworld", "tdmpc2", "jepa", "dreamerv3", "ppo", "sac"] | |
| ALGO_LABELS = { | |
| "hamworld": "HaM-World", | |
| "tdmpc2": "TD-MPC2", | |
| "jepa": "JEPA", | |
| "dreamerv3": "DreamerV3", | |
| "ppo": "PPO", | |
| "sac": "SAC", | |
| } | |
| ALGO_COLORS = { | |
| "hamworld": "#2563eb", | |
| "tdmpc2": "#d97706", | |
| "jepa": "#dc2626", | |
| "dreamerv3": "#059669", | |
| "ppo": "#7c3aed", | |
| "sac": "#0f766e", | |
| } | |
| TASK_ORDER = ["reacher_easy", "finger_spin", "cheetah_run", "cartpole_swingup"] | |
| TASK_LABELS = { | |
| "reacher_easy": "Reacher", | |
| "finger_spin": "Finger", | |
| "cheetah_run": "Cheetah", | |
| "cartpole_swingup": "Cartpole", | |
| } | |
| def parse_args() -> argparse.Namespace: | |
| parser = argparse.ArgumentParser(description="Analyze OOD evaluation results and export retention summaries/plots.") | |
| parser.add_argument("--input-csv", required=True, help="Path to ood_results.csv") | |
| parser.add_argument("--summary-dir", default=None, help="Directory for summary CSV outputs.") | |
| parser.add_argument("--figures-dir", default=None, help="Directory for figure outputs.") | |
| return parser.parse_args() | |
| def _to_float(value: str) -> float: | |
| try: | |
| return float(value) | |
| except ValueError: | |
| return float("nan") | |
| def _safe_mean(values: list[float]) -> float: | |
| finite = [value for value in values if not math.isnan(value)] | |
| if not finite: | |
| return float("nan") | |
| return statistics.fmean(finite) | |
| def _format_float(value: float) -> str: | |
| return "" if math.isnan(value) else f"{value:.6f}" | |
| def load_rows(path: Path) -> list[dict[str, object]]: | |
| rows: list[dict[str, object]] = [] | |
| with path.open("r", encoding="utf-8", newline="") as handle: | |
| for row in csv.DictReader(handle): | |
| rows.append( | |
| { | |
| "algo": row["algo"], | |
| "task": row["task"], | |
| "seed": int(float(row["seed"])), | |
| "condition": row["condition"], | |
| "episodes": int(float(row["episodes"])), | |
| "return_mean": _to_float(row["return_mean"]), | |
| "return_std": _to_float(row["return_std"]), | |
| "mse_horizon": int(float(row["mse_horizon"])), | |
| "mse_mean": _to_float(row["mse_mean"]), | |
| } | |
| ) | |
| return rows | |
| def summarize(rows: list[dict[str, object]]) -> tuple[list[dict[str, object]], list[dict[str, object]]]: | |
| seed_id_return = { | |
| (str(row["algo"]), str(row["task"]), int(row["seed"])): float(row["return_mean"]) | |
| for row in rows | |
| if str(row["condition"]) == "id" | |
| } | |
| seed_id_mse = { | |
| (str(row["algo"]), str(row["task"]), int(row["seed"])): float(row["mse_mean"]) | |
| for row in rows | |
| if str(row["condition"]) == "id" | |
| } | |
| condition_rows: list[dict[str, object]] = [] | |
| task_groups: dict[tuple[str, str], dict[str, object]] = {} | |
| condition_grouped: dict[tuple[str, str, str], list[dict[str, object]]] = defaultdict(list) | |
| for row in rows: | |
| if str(row["condition"]) == "id": | |
| continue | |
| condition_grouped[(str(row["algo"]), str(row["task"]), str(row["condition"]))].append(row) | |
| for (algo, task, condition), group in sorted(condition_grouped.items()): | |
| retention_values = [] | |
| mse_ratio_values = [] | |
| for row in group: | |
| key = (algo, task, int(row["seed"])) | |
| id_return = seed_id_return.get(key, float("nan")) | |
| id_mse = seed_id_mse.get(key, float("nan")) | |
| if abs(id_return) > 1e-9 and not math.isnan(id_return): | |
| retention_values.append(float(row["return_mean"]) / id_return) | |
| if not math.isnan(id_mse) and id_mse > 1e-12 and not math.isnan(float(row["mse_mean"])): | |
| mse_ratio_values.append(float(row["mse_mean"]) / id_mse) | |
| condition_row = { | |
| "algo": algo, | |
| "task": task, | |
| "condition": condition, | |
| "n": len(group), | |
| "id_return_mean": _safe_mean([seed_id_return.get((algo, task, int(row["seed"])), float("nan")) for row in group]), | |
| "ood_return_mean": _safe_mean([float(row["return_mean"]) for row in group]), | |
| "avg_retention": _safe_mean(retention_values), | |
| "mse_id_mean": _safe_mean([seed_id_mse.get((algo, task, int(row["seed"])), float("nan")) for row in group]), | |
| "mse_ood_mean": _safe_mean([float(row["mse_mean"]) for row in group]), | |
| "mse_ratio": _safe_mean(mse_ratio_values), | |
| } | |
| condition_rows.append(condition_row) | |
| task_group = task_groups.setdefault( | |
| (algo, task), | |
| { | |
| "algo": algo, | |
| "task": task, | |
| "retentions": [], | |
| "ood_returns": [], | |
| "id_returns": [], | |
| "mse_ids": [], | |
| "mse_oods": [], | |
| "condition_retentions": [], | |
| }, | |
| ) | |
| task_group["retentions"].extend(retention_values) | |
| task_group["ood_returns"].extend(float(row["return_mean"]) for row in group) | |
| task_group["id_returns"].extend( | |
| seed_id_return.get((algo, task, int(row["seed"])), float("nan")) for row in group | |
| ) | |
| task_group["mse_ids"].extend( | |
| seed_id_mse.get((algo, task, int(row["seed"])), float("nan")) for row in group | |
| ) | |
| task_group["mse_oods"].extend(float(row["mse_mean"]) for row in group) | |
| if not math.isnan(condition_row["avg_retention"]): | |
| task_group["condition_retentions"].append((condition, float(condition_row["avg_retention"]))) | |
| ranking_rows: list[dict[str, object]] = [] | |
| overall_rows: list[dict[str, object]] = [] | |
| for (algo, task), stats in sorted(task_groups.items()): | |
| avg_retention = _safe_mean(stats["retentions"]) | |
| avg_id_return = _safe_mean(stats["id_returns"]) | |
| avg_ood_return = _safe_mean(stats["ood_returns"]) | |
| avg_mse_id = _safe_mean(stats["mse_ids"]) | |
| avg_mse_ood = _safe_mean(stats["mse_oods"]) | |
| mse_ratio = float("nan") | |
| if not math.isnan(avg_mse_id) and avg_mse_id > 1e-12 and not math.isnan(avg_mse_ood): | |
| mse_ratio = avg_mse_ood / avg_mse_id | |
| worst_condition = "" | |
| worst_retention = float("nan") | |
| if stats["condition_retentions"]: | |
| worst_condition, worst_retention = min(stats["condition_retentions"], key=lambda item: item[1]) | |
| ranking_rows.append( | |
| { | |
| "scope": "task", | |
| "algo": algo, | |
| "task": task, | |
| "avg_id_return": avg_id_return, | |
| "avg_ood_return": avg_ood_return, | |
| "ood_delta": avg_ood_return - avg_id_return if not math.isnan(avg_id_return) and not math.isnan(avg_ood_return) else float("nan"), | |
| "avg_retention": avg_retention, | |
| "worst_condition": worst_condition, | |
| "worst_condition_retention": worst_retention, | |
| "avg_mse_id": avg_mse_id, | |
| "avg_mse_ood": avg_mse_ood, | |
| "mse_ratio": mse_ratio, | |
| "n_ood_rows": len(stats["ood_returns"]), | |
| } | |
| ) | |
| overall_rows.append( | |
| { | |
| "algo": algo, | |
| "task": task, | |
| "avg_retention": avg_retention, | |
| "worst_condition_retention": worst_retention, | |
| } | |
| ) | |
| overall_grouped: dict[str, list[dict[str, object]]] = defaultdict(list) | |
| for row in overall_rows: | |
| overall_grouped[str(row["algo"])].append(row) | |
| for algo, group in sorted(overall_grouped.items()): | |
| ranking_rows.append( | |
| { | |
| "scope": "overall", | |
| "algo": algo, | |
| "task": "", | |
| "avg_id_return": float("nan"), | |
| "avg_ood_return": float("nan"), | |
| "ood_delta": float("nan"), | |
| "avg_retention": _safe_mean([float(row["avg_retention"]) for row in group]), | |
| "worst_condition": "", | |
| "worst_condition_retention": _safe_mean([float(row["worst_condition_retention"]) for row in group]), | |
| "avg_mse_id": float("nan"), | |
| "avg_mse_ood": float("nan"), | |
| "mse_ratio": float("nan"), | |
| "n_ood_rows": sum(int(task_groups[(algo, str(row["task"]))]["ood_returns"].__len__()) for row in group), | |
| } | |
| ) | |
| def _rank(rows_in: list[dict[str, object]], key: str) -> None: | |
| sorted_rows = sorted(rows_in, key=lambda row: float(row[key]), reverse=True) | |
| for index, row in enumerate(sorted_rows, start=1): | |
| row[f"rank_by_{key}"] = index | |
| _rank([row for row in ranking_rows if row["scope"] == "overall"], "avg_retention") | |
| for task in sorted({str(row["task"]) for row in ranking_rows if row["scope"] == "task"}): | |
| _rank([row for row in ranking_rows if row["scope"] == "task" and row["task"] == task], "avg_retention") | |
| return ranking_rows, condition_rows | |
| def write_csv(path: Path, rows: list[dict[str, object]], fieldnames: list[str]) -> None: | |
| path.parent.mkdir(parents=True, exist_ok=True) | |
| with path.open("w", encoding="utf-8", newline="") as handle: | |
| writer = csv.DictWriter(handle, fieldnames=fieldnames) | |
| writer.writeheader() | |
| for row in rows: | |
| normalized = {} | |
| for key in fieldnames: | |
| value = row.get(key, "") | |
| if isinstance(value, float): | |
| normalized[key] = _format_float(value) | |
| else: | |
| normalized[key] = value | |
| writer.writerow(normalized) | |
| def _plot_metric(ranking_rows: list[dict[str, object]], metric_key: str, output_path: Path, title: str, y_label: str) -> None: | |
| task_rows = [row for row in ranking_rows if row["scope"] == "task"] | |
| tasks = [task for task in TASK_ORDER if any(str(row["task"]) == task for row in task_rows)] | |
| if not tasks: | |
| return | |
| ncols = 2 if len(tasks) > 1 else 1 | |
| nrows = (len(tasks) + ncols - 1) // ncols | |
| fig, axes = plt.subplots(nrows, ncols, figsize=(5.4 * ncols, 4.2 * nrows), dpi=200, sharey=True) | |
| fig.patch.set_facecolor("#f4f7fb") | |
| if hasattr(axes, "flatten"): | |
| axes = list(axes.flatten()) | |
| else: | |
| axes = [axes] | |
| for axis, task in zip(axes, tasks): | |
| rows = [row for row in task_rows if row["task"] == task] | |
| rows.sort(key=lambda row: ALGO_ORDER.index(str(row["algo"])) if str(row["algo"]) in ALGO_ORDER else 999) | |
| labels = [ALGO_LABELS.get(str(row["algo"]), str(row["algo"])) for row in rows] | |
| values = [float(row[metric_key]) for row in rows] | |
| colors = [ALGO_COLORS.get(str(row["algo"]), "#4b5563") for row in rows] | |
| bars = axis.bar(labels, values, color=colors, width=0.72) | |
| axis.axhline(1.0, color="#94a3b8", linewidth=1.0, linestyle="--") | |
| axis.set_title(TASK_LABELS.get(task, task), loc="left", fontsize=13, fontweight="bold", color="#111827") | |
| axis.set_ylabel(y_label, fontsize=11, color="#334155") | |
| axis.set_facecolor("#ffffff") | |
| axis.grid(True, axis="y", color="#e5e7eb", linewidth=1.0) | |
| axis.spines["top"].set_visible(False) | |
| axis.spines["right"].set_visible(False) | |
| axis.spines["left"].set_color("#94a3b8") | |
| axis.spines["bottom"].set_color("#94a3b8") | |
| axis.tick_params(axis="x", labelrotation=15, colors="#475569") | |
| axis.tick_params(axis="y", colors="#475569") | |
| for bar, value in zip(bars, values): | |
| axis.text( | |
| bar.get_x() + bar.get_width() * 0.5, | |
| value + 0.02 * max(1.0, max(values)), | |
| f"{value:.3f}", | |
| ha="center", | |
| va="bottom", | |
| fontsize=9, | |
| color="#334155", | |
| ) | |
| for axis in axes[len(tasks):]: | |
| axis.set_visible(False) | |
| fig.suptitle(title, x=0.06, y=0.99, ha="left", fontsize=18, fontweight="bold", color="#111827") | |
| fig.text(0.06, 0.91, "Dashed line marks retention=1.0; higher is better.", fontsize=10.5, color="#475569", ha="left") | |
| fig.subplots_adjust(left=0.08, right=0.98, bottom=0.12, top=0.88 if len(tasks) > 2 else 0.82, wspace=0.16, hspace=0.28) | |
| output_path.parent.mkdir(parents=True, exist_ok=True) | |
| fig.savefig(output_path, facecolor=fig.get_facecolor(), bbox_inches="tight") | |
| plt.close(fig) | |
| def main() -> int: | |
| args = parse_args() | |
| input_csv = Path(args.input_csv).expanduser().resolve() | |
| if not input_csv.exists(): | |
| raise SystemExit(f"Input CSV does not exist: {input_csv}") | |
| campaign_root = input_csv.parent.parent | |
| summary_dir = Path(args.summary_dir).expanduser().resolve() if args.summary_dir else campaign_root / "summary" | |
| figures_dir = Path(args.figures_dir).expanduser().resolve() if args.figures_dir else campaign_root / "figures" | |
| rows = load_rows(input_csv) | |
| ranking_rows, condition_rows = summarize(rows) | |
| ranking_fieldnames = [ | |
| "scope", | |
| "algo", | |
| "task", | |
| "rank_by_avg_retention", | |
| "avg_id_return", | |
| "avg_ood_return", | |
| "ood_delta", | |
| "avg_retention", | |
| "worst_condition", | |
| "worst_condition_retention", | |
| "avg_mse_id", | |
| "avg_mse_ood", | |
| "mse_ratio", | |
| "n_ood_rows", | |
| ] | |
| condition_fieldnames = [ | |
| "algo", | |
| "task", | |
| "condition", | |
| "n", | |
| "id_return_mean", | |
| "ood_return_mean", | |
| "avg_retention", | |
| "mse_id_mean", | |
| "mse_ood_mean", | |
| "mse_ratio", | |
| ] | |
| write_csv(summary_dir / "ood_retention_rankings.csv", ranking_rows, ranking_fieldnames) | |
| write_csv(summary_dir / "ood_condition_retention.csv", condition_rows, condition_fieldnames) | |
| _plot_metric( | |
| ranking_rows, | |
| metric_key="avg_retention", | |
| output_path=figures_dir / "ood_avg_retention_by_task.png", | |
| title="OOD Average Retention By Task", | |
| y_label="average retention", | |
| ) | |
| _plot_metric( | |
| ranking_rows, | |
| metric_key="worst_condition_retention", | |
| output_path=figures_dir / "ood_worst_retention_by_task.png", | |
| title="OOD Worst-Condition Retention By Task", | |
| y_label="worst-condition retention", | |
| ) | |
| print(summary_dir / "ood_retention_rankings.csv") | |
| print(summary_dir / "ood_condition_retention.csv") | |
| print(figures_dir / "ood_avg_retention_by_task.png") | |
| print(figures_dir / "ood_worst_retention_by_task.png") | |
| return 0 | |
| if __name__ == "__main__": | |
| raise SystemExit(main()) | |