Resume scripts + paper figures

- Fix EVAL_RE regex bug in adaptive_k_cellpolicy_pipeline.py so plateau-based
k-bumping actually fires (metric was always None before).
- Add --initial_adapter so adaptive-k variants can warm-start from a saved
LoRA adapter directory.
- Make strawman_cellpolicy_pipeline.sh accept an optional SFT_INIT env var
for warm-starting from a previously trained baseline checkpoint.
- New launch_finish_repos23.sh: brings up the 6 strawman warm-start
variants and 2 adaptive-k resume variants across all 8 GPUs.
- Add paper figures comparing latent vs vanilla baseline solve / per-cell
exact across the three curriculum stages (plot_stage_progression.py +
stage_progression_{solve,exact}.{pdf,png}).
- Update .gitignore to keep in-flight run output trees out of the code repo.

.gitattributes

@@ -35,3 +35,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
 *.jsonl filter=lfs diff=lfs merge=lfs -text
 *.log filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text
+*.pdf filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -25,3 +25,9 @@ tmp_latent_debug/
 logs/
 
 _pushlogs/
+
+# In-flight run output trees (kept in HF model repos, not the code repo)
+_runs/strawman_warm_*/
+_runs/adaptive_k_resume_*/
+_runs/launch_finish_repos23_*_pids.txt
+curriculum_cot/

_runs/_paper_figures/plot_stage_progression.py

@@ -0,0 +1,87 @@
+"""Paper-style figures: Solve rate + Per-cell exact across stages.  Two separate figures,
+no titles, no footer text — only axes, lines, markers, legend.
+"""
+from __future__ import annotations
+
+from pathlib import Path
+
+import matplotlib as mpl
+import matplotlib.pyplot as plt
+
+OUT_DIR = Path(__file__).resolve().parent
+
+# -------------------------------- DATA ---------------------------------------
+STAGES = ["Stage 1", "Stage 2", "Stage 3"]
+
+LATENT_SOLVE   = [0.70, 0.50, 0.58]
+LATENT_EXACT   = [0.95, 0.958, 0.967]
+
+BASELINE_SOLVE = [0.78, 0.40, 0.44]
+BASELINE_EXACT = [0.988, 0.88, 0.83]
+
+# -------------------------------- STYLE --------------------------------------
+mpl.rcParams.update({
+    "font.family": "serif",
+    "font.serif": ["DejaVu Serif", "Times New Roman", "Times", "Liberation Serif"],
+    "font.size": 12,
+    "axes.labelsize": 12,
+    "xtick.labelsize": 11,
+    "ytick.labelsize": 11,
+    "legend.fontsize": 11,
+    "axes.spines.top": False,
+    "axes.spines.right": False,
+    "axes.linewidth": 1.0,
+    "lines.linewidth": 2.0,
+    "lines.markersize": 8.5,
+    "xtick.direction": "in",
+    "ytick.direction": "in",
+    "xtick.major.size": 4,
+    "ytick.major.size": 4,
+    "pdf.fonttype": 42,
+    "ps.fonttype": 42,
+})
+
+LATENT_COLOR   = "#1f4f8b"
+BASELINE_COLOR = "#b21e2f"
+GRID_KW = dict(linestyle=":", linewidth=0.7, color="0.7", alpha=0.7)
+x = list(range(len(STAGES)))
+
+
+def _plot(y_latent, y_baseline, ylim, yticks, ylabel, fname):
+    fig, ax = plt.subplots(figsize=(4.6, 3.4), constrained_layout=True)
+    ax.plot(
+        x, y_latent,
+        color=LATENT_COLOR, marker="s", linestyle="-",
+        label="Latent (recurrent-hidden)",
+    )
+    ax.plot(
+        x, y_baseline,
+        color=BASELINE_COLOR, marker="o", linestyle="--",
+        label="Baseline (vanilla 1.5B)",
+    )
+    ax.set_xticks(x, STAGES)
+    ax.set_ylim(*ylim)
+    ax.set_yticks(yticks)
+    ax.set_ylabel(ylabel)
+    ax.grid(True, axis="y", **GRID_KW)
+    ax.legend(frameon=False, loc="best")
+    fig.savefig(OUT_DIR / f"{fname}.pdf", bbox_inches="tight")
+    fig.savefig(OUT_DIR / f"{fname}.png", dpi=300, bbox_inches="tight")
+    plt.close(fig)
+    print(f"saved {OUT_DIR / f'{fname}.pdf'}")
+
+
+_plot(
+    LATENT_SOLVE, BASELINE_SOLVE,
+    ylim=(0.0, 1.0),
+    yticks=[0.0, 0.2, 0.4, 0.6, 0.8, 1.0],
+    ylabel="Solve rate",
+    fname="stage_progression_solve",
+)
+_plot(
+    LATENT_EXACT, BASELINE_EXACT,
+    ylim=(0.78, 1.00),
+    yticks=[0.80, 0.84, 0.88, 0.92, 0.96, 1.00],
+    ylabel="Per-cell set-match rate",
+    fname="stage_progression_exact",
+)

_runs/_paper_figures/stage_progression_exact.pdf

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0734e025a134db1521ff9e06a7a1d00fac1a7786ed7ddd48aba25082b25cfa16
+size 12154

_runs/_paper_figures/stage_progression_solve.pdf

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:004adf9132094bc6d0c639c3cca706be383b77efde5ad5313bd1f26beb656f38
+size 11689

_runs/adaptive_k_cellpolicy_pipeline.py

@@ -76,16 +76,27 @@ def parse_args() -> argparse.Namespace:
     p.add_argument("--train_rows", type=int, default=10000)
     p.add_argument("--enable_gc", action="store_true", default=True)
     p.add_argument("--seed", type=int, default=0)
+    p.add_argument(
+        "--initial_adapter",
+        default="",
+        help="Optional adapter dir to warm-start the first SFT phase from (used for resuming an interrupted adaptive-k run at a non-zero k).",
+    )
     return p.parse_args()
 
 
 # ---- log parsing -----------------------------------------------------------
 
-EVAL_RE = re.compile(r"exact_set_match_rate.*?([01]\.\d+)")
+# The SFT trainer logs eval as e.g.
+#   [latent sft eval] parse=1.000 canonical=1.000 exact=0.685 precision=... recall=... solve=0.000
+# i.e. the field is shortened from `exact_set_match_rate` to `exact=`. The
+# original regex `exact_set_match_rate.*?([01]\.\d+)` therefore never matched,
+# leaving `eval_exact_set_match_rate=None` in PIPELINE.log and disabling
+# plateau-based k-bumping. Fix: parse the short form on `[* sft eval]` lines.
+EVAL_RE = re.compile(r"\[\w+ sft eval\][^\n]*?exact=([01]\.\d+)")
 
 
 def latest_eval_metric(log_path: Path) -> Optional[float]:
-    """Return the most recent eval exact_set_match_rate from the SFT train log."""
+    """Return the most recent eval per-cell exact match rate from the SFT train log."""
     if not log_path.exists():
         return None
     last: Optional[float] = None
@@ -343,7 +354,9 @@ def main() -> None:
 
     history: List[dict] = []
     cur_k = int(args.start_k)
-    cur_init: str = ""  # "" -> train from base
+    cur_init: str = args.initial_adapter or ""
+    if cur_init:
+        log(f"  initial_adapter={cur_init}")
     last_metric_at_k: Optional[float] = None
     phases_at_k = 0
     sft_phase_idx = 0

_runs/launch_finish_repos23.sh

@@ -0,0 +1,132 @@
+#!/usr/bin/env bash
+# Finish repos 2 (strawman) and 3 (adaptive-k) on all 8 GPUs.
+#
+# GPU 0,1: resume the two half-finished adaptive-k runs at the next k
+#          (init from the latest HF checkpoint of phase02_k0 / phase01_k0).
+# GPU 2-7: six warm-started strawman variants, all initialised from the
+#          v6_i_sft_v_oversample10 best checkpoint (solve=0.44 @ step 200),
+#          sweeping the SFT LR / multi-value oversample / GRPO reward axes.
+#
+# After all eight jobs exit we will have:
+#   - adaptive-k variants pushed past their phase-2 plateau
+#   - 6 strawman variants whose SFT loss is guaranteed to decrease (warm init)
+#     plus a final GRPO phase with a 100-puzzle eval, so each variant publishes
+#     a paper-ready (per-cell exact, solve) tuple.
+set -euo pipefail
+
+ROOT="/home/ubuntu/curriculum_cot"
+HFROOT="/home/ubuntu/hf_checkpoints"
+TS="$(date +%Y%m%d_%H%M%S)"
+ADAPT_SWEEP="${ROOT}/_runs/adaptive_k_resume_${TS}"
+STRAW_SWEEP="${ROOT}/_runs/strawman_warm_${TS}"
+mkdir -p "${ADAPT_SWEEP}" "${STRAW_SWEEP}"
+PY="${ROOT}/_runs/adaptive_k_cellpolicy_pipeline.py"
+STRAW_PIPE="${ROOT}/_runs/strawman_cellpolicy_pipeline.sh"
+chmod +x "${STRAW_PIPE}"
+
+# Warm-start checkpoint for strawman: v6_i_sft_v_oversample10 S3 SFT step 200
+# (best baseline anywhere in the sweep: per-cell exact 0.951, solve 0.440).
+V6I_BEST="${HFROOT}/baseline/v6_i_sft_v_oversample10/s3_sft/checkpoint-step-00200"
+if [[ ! -f "${V6I_BEST}/adapter_model.safetensors" ]]; then
+  echo "ERROR: missing v6_i warm-start checkpoint at ${V6I_BEST}" >&2
+  exit 1
+fi
+
+PIDS_FILE="${ADAPT_SWEEP}/../launch_finish_repos23_${TS}_pids.txt"
+: > "${PIDS_FILE}"
+
+launch_adaptive() {
+  local variant="$1" gpu="$2" init_adapter="$3" start_k="$4" max_k="$5" eps="$6" max_phases="$7"
+  local out="${ADAPT_SWEEP}/${variant}"
+  mkdir -p "${out}"
+  echo "[launch] ADAPTIVE ${variant} on GPU ${gpu}  out=${out}  init=${init_adapter}"
+  nohup /opt/pytorch/bin/python -u "${PY}" \
+    --variant "${variant}" \
+    --gpu "${gpu}" \
+    --output_root "${out}" \
+    --initial_adapter "${init_adapter}" \
+    --start_k "${start_k}" \
+    --max_k "${max_k}" \
+    --steps_per_phase 600 \
+    --max_phases_per_k "${max_phases}" \
+    --plateau_eps "${eps}" \
+    --sft_lr 2e-5 --sft_bs 8 --sft_ga 4 \
+    --grpo_steps 1500 --grpo_lr 5e-6 --grpo_bs 8 --grpo_ga 4 --grpo_ng 8 \
+    > "${out}/console.log" 2>&1 &
+  local pid=$!
+  disown "${pid}" || true
+  echo "ADAPTIVE_${variant}=${pid}" >> "${PIDS_FILE}"
+}
+
+launch_strawman() {
+  # Usage: launch_strawman <variant> <gpu> <KEY=VALUE>...
+  local variant="$1" gpu="$2"
+  shift 2
+  local out="${STRAW_SWEEP}/${variant}"
+  mkdir -p "${out}"
+  echo "[launch] STRAWMAN ${variant} on GPU ${gpu}  out=${out}"
+  # NB: env KEY=VAL ... bash PIPE   — keep all assignments before `bash` and
+  # use printf '%q' so values like `2e-5` round-trip through env quoting safely.
+  local assignments=""
+  for kv in "$@"; do
+    assignments+=" ${kv}"
+  done
+  nohup env VARIANT="${variant}" GPU="${gpu}" OUTPUT_ROOT="${out}" \
+        SFT_INIT="${V6I_BEST}" \
+        ${assignments} \
+        bash "${STRAW_PIPE}" > "${out}/console.log" 2>&1 &
+  local pid=$!
+  disown "${pid}" || true
+  echo "STRAWMAN_${variant}=${pid}" >> "${PIDS_FILE}"
+}
+
+# === Adaptive-k resumes (GPUs 0,1) ===========================================
+# Pick the latest k=0 ckpt for each, resume at k=1 with a stricter plateau and
+# 3 phases per k so the bumping logic actually has a chance to fire now that
+# the regex bug is fixed.
+ADAPT_A_INIT="${HFROOT}/adaptive_k/20260525_024629/adaptive_a_eps01/sft_phase02_k0/checkpoint-step-00600"
+ADAPT_B_INIT="${HFROOT}/adaptive_k/20260525_024629/adaptive_b_fastgrow/sft_phase01_k0/checkpoint-step-00800"
+
+launch_adaptive adaptive_e_resume_a 0 "${ADAPT_A_INIT}" 1 4 0.005 3
+launch_adaptive adaptive_f_resume_b 1 "${ADAPT_B_INIT}" 1 4 0.005 3
+
+# === Strawman warm-started from v6_i (GPUs 2-7) =============================
+# Six variants spanning SFT LR x oversample x GRPO reward to make sure at least
+# one converges robustly. All start from the same warm checkpoint so the SFT
+# loss is guaranteed to decrease from the first step.
+
+launch_strawman strawman_warm_a_lr2e5_o3   2 \
+  SFT_LR=2e-5 SFT_OVERSAMPLE=3 SFT_MAX_STEPS=1500 \
+  GRPO_LR=5e-6 GRPO_MAX_STEPS=1500 \
+  EXACT_MATCH_BONUS=1.0 CARD_MISMATCH_PEN=1.5
+
+launch_strawman strawman_warm_b_lr2e5_o5   3 \
+  SFT_LR=2e-5 SFT_OVERSAMPLE=5 SFT_MAX_STEPS=1500 \
+  GRPO_LR=5e-6 GRPO_MAX_STEPS=1500 \
+  EXACT_MATCH_BONUS=1.0 CARD_MISMATCH_PEN=1.5
+
+launch_strawman strawman_warm_c_lr2e5_o8   4 \
+  SFT_LR=2e-5 SFT_OVERSAMPLE=8 SFT_MAX_STEPS=1500 \
+  GRPO_LR=5e-6 GRPO_MAX_STEPS=1500 \
+  EXACT_MATCH_BONUS=1.0 CARD_MISMATCH_PEN=1.5
+
+launch_strawman strawman_warm_d_lr5e5_o10  5 \
+  SFT_LR=5e-5 SFT_OVERSAMPLE=10 SFT_MAX_STEPS=1500 \
+  GRPO_LR=5e-6 GRPO_MAX_STEPS=1500 \
+  EXACT_MATCH_BONUS=1.0 CARD_MISMATCH_PEN=1.5
+
+launch_strawman strawman_warm_e_lr1e5_o5   6 \
+  SFT_LR=1e-5 SFT_OVERSAMPLE=5 SFT_MAX_STEPS=1500 \
+  GRPO_LR=2e-6 GRPO_MAX_STEPS=1500 \
+  EXACT_MATCH_BONUS=1.0 CARD_MISMATCH_PEN=1.5
+
+launch_strawman strawman_warm_f_lr2e5_o5_sharp 7 \
+  SFT_LR=2e-5 SFT_OVERSAMPLE=5 SFT_MAX_STEPS=1500 \
+  GRPO_LR=2e-6 GRPO_MAX_STEPS=1500 \
+  EXACT_MATCH_BONUS=4.0 CARD_MISMATCH_PEN=2.0
+
+echo
+echo "[launch] adaptive sweep root: ${ADAPT_SWEEP}"
+echo "[launch] strawman sweep root: ${STRAW_SWEEP}"
+echo "[launch] PIDs file: ${PIDS_FILE}"
+cat "${PIDS_FILE}"

_runs/strawman_cellpolicy_pipeline.sh

@@ -54,6 +54,10 @@ EVAL_ROWS="${EVAL_ROWS:-100}"
 TRAIN_ROWS="${TRAIN_ROWS:-10000}"
 USE_GC="${USE_GC:-1}"  # GC=1 to allow bs 16 on a single 80G GPU
 PHASE_WALL_SECS="${PHASE_WALL_SECS:-0}"
+# Optional warm-start. Empty = train SFT from base Qwen (true strawman).
+# Otherwise must point at a directory containing an `adapter_model.safetensors`
+# (e.g. a previously trained S3 SFT checkpoint).
+SFT_INIT="${SFT_INIT:-}"
 
 TRAIN_JSONL="${ROOT}/data/sudoku_t3_20empty_value_qwen_text_stage1_train.jsonl"
 EVAL_JSONL="${ROOT}/data/sudoku_t3_20empty_value_qwen_text_stage1_eval.jsonl"
@@ -90,6 +94,7 @@ if [[ "${USE_GC}" == "1" ]]; then GC_FLAG=(--enable_gradient_checkpointing); fi
 
 log "===== STRAWMAN ${VARIANT} on GPU ${GPU} ====="
 log "    SFT lr=${SFT_LR} max_steps=${SFT_MAX_STEPS}  bs=${SFT_BS}x${SFT_GA}  GC=${USE_GC}"
+log "    SFT_INIT=${SFT_INIT:-<base>}"
 log "    GRPO lr=${GRPO_LR} max_steps=${GRPO_MAX_STEPS} ng=${GRPO_NG}  bs=${GRPO_BS}x${GRPO_GA}"
 log "    rewards good=${REWARD_GOOD} bad=${PENALTY_BAD} mal=${PENALTY_MAL} empty=${PENALTY_EMPTY} sng=${PENALTY_SINGLETON} miss=${PENALTY_MISSING} bonus=${EXACT_MATCH_BONUS} card=${CARD_MISMATCH_PEN}"
 log "    out=${OUTPUT_ROOT}"
@@ -97,14 +102,14 @@ log "    out=${OUTPUT_ROOT}"
 # ----- Phase 1: SFT at stage_i=3 from BASE (no init adapter) -----
 SFT_DIR="${OUTPUT_ROOT}/sft"
 mkdir -p "${SFT_DIR}"
-log "=== PHASE SFT (stage_i=3, init=BASE) ==="
+log "=== PHASE SFT (stage_i=3, init=${SFT_INIT:-BASE}) ==="
 "${PYTHON_BIN}" -u "${SFT_SCRIPT}" \
   --model_name "${MODEL_NAME}" \
   --train_jsonl "${TRAIN_JSONL}" \
   --eval_jsonl "${EVAL_JSONL}" \
   --output_dir "${SFT_DIR}" \
   --cache_dir "${ROOT}/.hf_cache" \
-  --init_adapter_dir "" \
+  --init_adapter_dir "${SFT_INIT}" \
   --seed 0 \
   --gpu_id 0 \
   --stage_i 3 \