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_ckpt_evolution.py | 105 +++++++++++++++++++++++++++++++ 1 file changed, 105 insertions(+) create mode 100644 research/flossing/plot_ckpt_evolution.py (limited to 'research/flossing/plot_ckpt_evolution.py') diff --git a/research/flossing/plot_ckpt_evolution.py b/research/flossing/plot_ckpt_evolution.py new file mode 100644 index 0000000..1bde7b4 --- /dev/null +++ b/research/flossing/plot_ckpt_evolution.py @@ -0,0 +1,105 @@ +"""Plot λ-spectrum and accuracy evolution across training checkpoints.""" +import numpy as np +import matplotlib +matplotlib.use("Agg") +import matplotlib.pyplot as plt + +CKPT_STEPS = [2604, 7812, 13020, 18228, 20832, 26040] +CKPT_EPOCHS = [int(s * 20000 / 26041) for s in CKPT_STEPS] # actual epoch (rounded) +ROOT = "/home/yurenh2/rrm/research/flossing/ckpt_evolution" +OUT_DIR = "/home/yurenh2/rrm/research/flossing/plots_evolution" +import os; os.makedirs(OUT_DIR, exist_ok=True) + +datas = {} +for s in CKPT_STEPS: + d = np.load(f"{ROOT}/step_{s}.npz") + datas[s] = d + +# --- plot 1: λ_1 vs training step, split by success/failure --- +fig, axes = plt.subplots(1, 2, figsize=(13, 5)) + +ax = axes[0] +lam1_s_mean, lam1_f_mean = [], [] +lam1_s_std, lam1_f_std = [], [] +mean_s, mean_f = [], [] +mean_s_std, mean_f_std = [], [] +accs = [] +for s in CKPT_STEPS: + d = datas[s] + succ = d["exact_correct"] > 0.5 + L = d["lyap_spec"] + lam1_s_mean.append(L[succ,0].mean()); lam1_s_std.append(L[succ,0].std()) + lam1_f_mean.append(L[~succ,0].mean()); lam1_f_std.append(L[~succ,0].std()) + mean_s.append(L[succ].mean(axis=1).mean()); mean_s_std.append(L[succ].mean(axis=1).std()) + mean_f.append(L[~succ].mean(axis=1).mean()); mean_f_std.append(L[~succ].mean(axis=1).std()) + accs.append(succ.mean()) + +x = CKPT_EPOCHS +ax.errorbar(x, lam1_s_mean, yerr=lam1_s_std, fmt="C0-o", capsize=4, label=r"$\lambda_1$ success") +ax.errorbar(x, lam1_f_mean, yerr=lam1_f_std, fmt="C3-o", capsize=4, label=r"$\lambda_1$ failure") +ax.errorbar(x, mean_s, yerr=mean_s_std, fmt="C0--s", alpha=0.6, capsize=3, label=r"mean$_{1:8}$ success") +ax.errorbar(x, mean_f, yerr=mean_f_std, fmt="C3--s", alpha=0.6, capsize=3, label=r"mean$_{1:8}$ failure") +ax.axhline(0, color="k", ls=":", lw=0.6) +ax.set_xlabel("training epoch") +ax.set_ylabel(r"Lyapunov exponent (per inner cycle)") +ax.set_title("λ evolution during training\n(error bars = within-group std)") +ax.legend(loc="best", fontsize=9) +ax.grid(alpha=0.3) + +ax2 = axes[1] +gap = np.array(lam1_f_mean) - np.array(lam1_s_mean) +ax2.plot(x, gap, "C2-o", label=r"$\lambda_1^{fail} - \lambda_1^{succ}$", lw=2) +ax2.axhline(0, color="k", ls=":", lw=0.6) +ax2b = ax2.twinx() +ax2b.plot(x, accs, "C4-^", label="exact accuracy") +ax2b.set_ylabel("test exact accuracy", color="C4") +ax2b.tick_params(axis="y", labelcolor="C4") +ax2.set_xlabel("training epoch") +ax2.set_ylabel("λ gap (fail − succ)", color="C2") +ax2.tick_params(axis="y", labelcolor="C2") +ax2.set_title("succ/fail λ gap vs accuracy") +ax2.legend(loc="upper left") +ax2b.legend(loc="lower right") +ax2.grid(alpha=0.3) + +fig.tight_layout() +fig.savefig(f"{OUT_DIR}/lyap_vs_training.png", dpi=130) +plt.close() + +# --- plot 2: per-checkpoint λ spectrum (all 8) shape --- +fig, axes = plt.subplots(2, 3, figsize=(15, 7), sharey=True) +for ax, s in zip(axes.flat, CKPT_STEPS): + d = datas[s] + succ = d["exact_correct"] > 0.5 + L = d["lyap_spec"] + idx = np.arange(1, L.shape[1]+1) + mean_s_arr = L[succ].mean(0); std_s_arr = L[succ].std(0) + mean_f_arr = L[~succ].mean(0); std_f_arr = L[~succ].std(0) + ax.fill_between(idx, mean_s_arr-std_s_arr, mean_s_arr+std_s_arr, color="C0", alpha=0.2) + ax.plot(idx, mean_s_arr, "C0-o", label=f"succ n={succ.sum()}") + ax.fill_between(idx, mean_f_arr-std_f_arr, mean_f_arr+std_f_arr, color="C3", alpha=0.2) + ax.plot(idx, mean_f_arr, "C3-o", label=f"fail n={(~succ).sum()}") + ax.axhline(0, color="k", ls=":", lw=0.6) + ax.set_title(f"step_{s} (epoch≈{int(s*20000/26041)}, acc={succ.mean():.2%})") + ax.set_xlabel("λ index") + ax.legend(fontsize=8, loc="best") + ax.grid(alpha=0.3) +axes[0,0].set_ylabel(r"$\lambda_i$") +axes[1,0].set_ylabel(r"$\lambda_i$") +fig.tight_layout() +fig.savefig(f"{OUT_DIR}/spectrum_per_ckpt.png", dpi=130) +plt.close() + +# --- summary table --- +print(f"{'epoch':>6} {'step':>6} {'acc':>7} {'λ_1(s)':>9} {'λ_1(f)':>9} {'gap':>7} " + f"{'mean(s)':>9} {'mean(f)':>9} {'mean_gap':>9}") +for s, ep in zip(CKPT_STEPS, CKPT_EPOCHS): + d = datas[s] + succ = d["exact_correct"] > 0.5 + L = d["lyap_spec"] + l1s, l1f = L[succ,0].mean(), L[~succ,0].mean() + ms, mf = L[succ].mean(), L[~succ].mean() + print(f"{ep:>6} {s:>6} {succ.mean():>7.2%} {l1s:+9.3f} {l1f:+9.3f} {l1f-l1s:+7.3f} " + f"{ms:+9.3f} {mf:+9.3f} {mf-ms:+9.3f}") + +print(f"\nplots → {OUT_DIR}/") -- cgit v1.2.3