rrm workspace: TRM/HRM/SRM code, Maze dataset, dynamical-analysis pipelineHEADmain

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 <noreply@anthropic.com>

1 files changed, 76 insertions, 0 deletions

diff --git a/research/flossing/analyze_step3.py b/research/flossing/analyze_step3.py
new file mode 100644
index 0000000..825408c
--- /dev/null
+++ b/research/flossing/analyze_step3.py
@@ -0,0 +1,76 @@
+"""Plot Step 3 A/B/C trajectories: λ and acc over training."""
+import json, os
+import numpy as np
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+
+ROOT = "/home/yurenh2/rrm/research/flossing"
+OUT = f"{ROOT}/plots_step3"
+os.makedirs(OUT, exist_ok=True)
+
+runs = {
+    "A: baseline (α=0)\nfrom step_18228": "step3_A_baseline_18228.json",
+    "B: CF α=10 λ*=-0.15\nfrom step_18228": "step3_B_rf_18228.json",
+    "C: CF α=10 λ*=-0.05\nfrom step_26040": "step3_C_rf_26040.json",
+}
+colors = {"A": "C0", "B": "C3", "C": "C2"}
+
+fig, axes = plt.subplots(2, 2, figsize=(13, 9))
+
+for label, fn in runs.items():
+    d = json.loads(open(f"{ROOT}/{fn}").read())
+    key = label.split(":")[0]
+    steps = [r["step"] for r in d["steps"]]
+    sup = np.array([r["sup_loss"] for r in d["steps"]])
+    rf = np.array([r["rf_loss"] for r in d["steps"]])
+    lyap_mean = np.array([r["lyap1_mean"] for r in d["steps"]])
+    lyap_max = np.array([r["lyap1_max"] for r in d["steps"]])
+    frac = np.array([r["frac_above_star"] for r in d["steps"]])
+
+    # Smooth (moving average)
+    def smooth(x, w=20):
+        if len(x) < w: return x
+        kernel = np.ones(w) / w
+        return np.convolve(x, kernel, mode="same")
+
+    # acc evals
+    eval_steps = [e["step"] for e in d["evals"]]
+    eval_accs = [e["acc"] for e in d["evals"]]
+
+    axes[0,0].plot(eval_steps, eval_accs, f"{colors[key]}-o", label=label)
+    axes[0,1].plot(steps, smooth(sup), f"{colors[key]}-", label=label, alpha=0.8)
+    axes[1,0].plot(steps, smooth(lyap_mean), f"{colors[key]}-", label=f"{key} mean", alpha=0.8)
+    axes[1,0].plot(steps, smooth(lyap_max), f"{colors[key]}--", label=f"{key} max", alpha=0.4)
+    axes[1,1].plot(steps, smooth(frac), f"{colors[key]}-", label=label, alpha=0.8)
+
+axes[0,0].set_title("Test exact accuracy vs training step")
+axes[0,0].set_xlabel("step"); axes[0,0].set_ylabel("exact_acc"); axes[0,0].legend(fontsize=8, loc="best"); axes[0,0].grid(alpha=0.3)
+
+axes[0,1].set_title("Supervised loss (smoothed, w=20)")
+axes[0,1].set_xlabel("step"); axes[0,1].set_ylabel("sup_loss"); axes[0,1].legend(fontsize=8); axes[0,1].grid(alpha=0.3)
+
+axes[1,0].set_title("λ_joint_1 trajectory (smoothed)")
+axes[1,0].axhline(0, color="k", ls=":", lw=0.6)
+axes[1,0].axhline(-0.05, color="C2", ls="--", lw=0.5, label="C λ*")
+axes[1,0].axhline(-0.15, color="C3", ls="--", lw=0.5, label="B λ*")
+axes[1,0].set_xlabel("step"); axes[1,0].set_ylabel(r"$\lambda_{joint,1}$"); axes[1,0].legend(fontsize=7); axes[1,0].grid(alpha=0.3)
+
+axes[1,1].set_title(r"fraction of batch samples with $\lambda > \lambda^*$")
+axes[1,1].set_xlabel("step"); axes[1,1].set_ylabel("fraction"); axes[1,1].legend(fontsize=8); axes[1,1].grid(alpha=0.3)
+
+fig.suptitle("Step 3: contractive flossing as training-time regularizer on HRM")
+fig.tight_layout()
+fig.savefig(f"{OUT}/step3_trajectories.png", dpi=130)
+plt.close()
+
+print(f"\n== summary ==")
+for label, fn in runs.items():
+    d = json.loads(open(f"{ROOT}/{fn}").read())
+    init = d["initial_acc"]; fin = d["final_acc"]
+    last_lyap = d["steps"][-1]
+    print(f"  {label.split(':')[0]}: acc {init:.3f} → {fin:.3f} (Δ {fin-init:+.4f})  "
+          f"final λ_mean={last_lyap['lyap1_mean']:+.3f} max={last_lyap['lyap1_max']:+.3f}  "
+          f"final frac>λ*={last_lyap['frac_above_star']:.2f}")
+
+print(f"\nplots → {OUT}/")