From 66e0d8b9fd4d0f7a2231d689c055e26fdf1cf04a Mon Sep 17 00:00:00 2001 From: YurenHao0426 Date: Sat, 13 Jun 2026 12:35:36 -0500 Subject: rrm workspace: TRM/HRM/SRM code, Maze dataset, dynamical-analysis pipeline Curated export for clone-and-run Maze training (2x A6000) + diagnostics. trm/hrm pretrain.py carry trajectory-augmentation code (backward-compatible). Heavy artifacts (checkpoints/wandb/npz) gitignored; see PROVENANCE.md. Co-Authored-By: Claude Fable 5 --- research/flossing/plot_late_perturb_robustness.py | 140 ++++++++++++++++++++++ 1 file changed, 140 insertions(+) create mode 100644 research/flossing/plot_late_perturb_robustness.py (limited to 'research/flossing/plot_late_perturb_robustness.py') diff --git a/research/flossing/plot_late_perturb_robustness.py b/research/flossing/plot_late_perturb_robustness.py new file mode 100644 index 0000000..d28b3f0 --- /dev/null +++ b/research/flossing/plot_late_perturb_robustness.py @@ -0,0 +1,140 @@ +"""Plot late recurrent-state perturbation robustness curves.""" +from __future__ import annotations + +import argparse +import csv +from pathlib import Path + +import matplotlib.pyplot as plt + + +def read_rows(paths: list[Path]) -> list[dict[str, str]]: + rows: list[dict[str, str]] = [] + for path in paths: + with path.open() as f: + rows.extend(csv.DictReader(f)) + return rows + + +def f(row: dict[str, str], key: str) -> float: + return float(row[key]) + + +def write_combined(path: Path, rows: list[dict[str, str]]) -> None: + keys: list[str] = [] + for row in rows: + for key in row: + if key not in keys: + keys.append(key) + with path.open("w", newline="") as out: + writer = csv.DictWriter(out, fieldnames=keys) + writer.writeheader() + writer.writerows(rows) + + +COLORS = { + "trm_baseline_best": "#334155", + "trm_multi4_best": "#0f766e", + "trm_multi4_final": "#dc2626", +} +MARKERS = { + "trm_baseline_best": "o", + "trm_multi4_best": "s", + "trm_multi4_final": "X", +} + + +def plot_metric_grid(rows: list[dict[str, str]], metric: str, ylabel: str, out: Path) -> None: + labels = sorted({r["label"] for r in rows}) + afters = sorted({int(float(r["perturb_after"])) for r in rows}) + fig, axes = plt.subplots(1, len(afters), figsize=(4.0 * len(afters), 4.2), sharey=True) + if len(afters) == 1: + axes = [axes] + for ax, after in zip(axes, afters): + for label in labels: + lr = [r for r in rows if r["label"] == label and int(float(r["perturb_after"])) == after] + lr.sort(key=lambda r: f(r, "sigma")) + xs = [f(r, "sigma") for r in lr] + ys = [f(r, metric) for r in lr] + ax.plot( + xs, + ys, + marker=MARKERS.get(label, "o"), + linewidth=2.0, + markersize=5, + color=COLORS.get(label), + label=label.replace("trm_", "").replace("_", " "), + ) + ax.set_xscale("symlog", linthresh=1e-3) + ax.set_title(f"after {after}") + ax.set_xlabel("σ") + ax.grid(alpha=0.23) + ax.set_ylim(-0.02, 1.02) + axes[0].set_ylabel(ylabel) + axes[-1].legend(frameon=False, fontsize=8, loc="best") + fig.suptitle(f"Late-state perturbation robustness: {ylabel}") + fig.tight_layout() + fig.savefig(out, dpi=220, bbox_inches="tight") + plt.close(fig) + + +def plot_sigma_slice(rows: list[dict[str, str]], sigma: float, out: Path) -> None: + labels = sorted({r["label"] for r in rows}) + metrics = [ + ("retain_mean_on_clean_success", "Clean-success retention"), + ("rescue_mean_on_clean_fail", "Clean-fail rescue"), + ] + fig, axes = plt.subplots(1, 2, figsize=(11.2, 4.6), sharex=True) + for ax, (metric, title) in zip(axes, metrics): + for label in labels: + lr = [ + r + for r in rows + if r["label"] == label and abs(float(r["sigma"]) - sigma) <= max(1e-12, sigma * 1e-6) + ] + lr.sort(key=lambda r: int(float(r["perturb_after"]))) + xs = [int(float(r["perturb_after"])) for r in lr] + ys = [f(r, metric) for r in lr] + ax.plot( + xs, + ys, + marker=MARKERS.get(label, "o"), + linewidth=2.1, + markersize=6, + color=COLORS.get(label), + label=label.replace("trm_", "").replace("_", " "), + ) + ax.set_title(title) + ax.set_xlabel("Perturb after ACT step") + ax.set_ylim(-0.02, 1.02) + ax.grid(alpha=0.23) + axes[0].set_ylabel("Conditional probability") + axes[1].legend(frameon=False, loc="best") + fig.suptitle(f"Late perturbation conditional behavior at σ={sigma:g}") + fig.tight_layout() + fig.savefig(out, dpi=220, bbox_inches="tight") + plt.close(fig) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--summaries", nargs="+", required=True) + parser.add_argument("--out-dir", required=True) + parser.add_argument("--slice-sigma", type=float, default=0.1) + args = parser.parse_args() + + out_dir = Path(args.out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + rows = read_rows([Path(p) for p in args.summaries]) + write_combined(out_dir / "late_perturb_robustness_combined.csv", rows) + plot_metric_grid(rows, "mean_rollout_exact", "Mean perturbed-rollout exact", out_dir / "late_perturb_mean_rollout_exact_grid.png") + plot_metric_grid(rows, "pass_at_k", "Pass@K exact", out_dir / "late_perturb_pass_at_k_grid.png") + plot_metric_grid(rows, "all_k", "All-K exact", out_dir / "late_perturb_all_k_grid.png") + plot_metric_grid(rows, "retain_mean_on_clean_success", "Clean-success retention", out_dir / "late_perturb_retention_grid.png") + plot_metric_grid(rows, "rescue_mean_on_clean_fail", "Clean-fail rescue", out_dir / "late_perturb_rescue_grid.png") + plot_sigma_slice(rows, args.slice_sigma, out_dir / f"late_perturb_retention_rescue_sigma{args.slice_sigma:g}.png") + print(f"wrote {out_dir}") + + +if __name__ == "__main__": + main() -- cgit v1.2.3