Add data/ JSONLs + _runs/ launch scripts (override .gitignore)

_runs/LATENT_PID.txt

@@ -0,0 +1 @@
+164065 0-7 latent_reproduction_20260524_062728

_runs/adaptive_k_cellpolicy_pipeline.py

@@ -0,0 +1,430 @@
+#!/usr/bin/env python3
+"""Adaptive-k cell-policy pipeline (no curriculum).
+
+Wraps the existing per-cell trainers to implement an "adaptive-k" schedule:
+the model is trained at stage_i=3 only (no curriculum), with the number of
+recurrent-hidden thought tokens k starting at 0 (vanilla SFT) and being
+incremented whenever the eval exact_set_match metric plateaus. Each phase
+runs ``sft_latent_multi_output_train.py`` for ``steps_per_phase`` SFT steps
+at fixed k, initialised from the previous phase's best checkpoint (so the
+recurrent-hidden bank persists). After the final SFT phase, ``grpo_residual_projector_latent_train.py``
+is invoked at the converged k.
+
+The trainer scripts, prompt template, and scoring function are the *same*
+ones used by every cell-policy / latent experiment. The only knob this
+orchestrator provides is the k-schedule; per-cell prompt+supervision is
+handled by the existing trainers.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import re
+import shutil
+import subprocess
+import sys
+import time
+from pathlib import Path
+from typing import List, Optional
+
+ROOT = Path(__file__).resolve().parent.parent
+SFT_SCRIPT = ROOT / "latent_multi_output_cell_policy" / "sft_latent_multi_output_train.py"
+GRPO_SCRIPT = ROOT / "latent_multi_output_cell_policy" / "grpo_residual_projector_latent_train.py"
+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"
+
+
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser()
+    p.add_argument("--variant", required=True)
+    p.add_argument("--gpu", required=True)
+    p.add_argument("--output_root", required=True)
+    p.add_argument("--model_name", default="Qwen/Qwen2.5-1.5B-Instruct")
+    p.add_argument("--cache_dir", default=str(ROOT / ".hf_cache"))
+    p.add_argument("--python_bin", default="/opt/pytorch/bin/python")
+    p.add_argument("--latent_mode", default="recurrent_hidden")
+    p.add_argument("--start_k", type=int, default=0)
+    p.add_argument("--max_k", type=int, default=4)
+    p.add_argument("--steps_per_phase", type=int, default=600)
+    p.add_argument(
+        "--max_phases_per_k",
+        type=int,
+        default=2,
+        help="Hard cap on how many ``steps_per_phase`` chunks to spend at a single k before bumping.",
+    )
+    p.add_argument(
+        "--plateau_eps",
+        type=float,
+        default=0.01,
+        help="If eval exact_set_match_rate improves by less than this between two consecutive phases at the same k, declare a plateau and bump k.",
+    )
+    p.add_argument("--sft_lr", type=float, default=2e-5)
+    p.add_argument("--sft_bs", type=int, default=8)
+    p.add_argument("--sft_ga", type=int, default=4)
+    p.add_argument("--sft_oversample", type=int, default=3)
+    p.add_argument("--grpo_steps", type=int, default=1500)
+    p.add_argument("--grpo_lr", type=float, default=5e-6)
+    p.add_argument("--grpo_bs", type=int, default=8)
+    p.add_argument("--grpo_ga", type=int, default=4)
+    p.add_argument("--grpo_ng", type=int, default=8)
+    p.add_argument("--grpo_beta", type=float, default=0.0)
+    p.add_argument("--grpo_max_prompt", type=int, default=768)
+    p.add_argument("--grpo_max_completion", type=int, default=24)
+    p.add_argument("--eval_rows", type=int, default=100)
+    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)
+    return p.parse_args()
+
+
+# ---- log parsing -----------------------------------------------------------
+
+EVAL_RE = re.compile(r"exact_set_match_rate.*?([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."""
+    if not log_path.exists():
+        return None
+    last: Optional[float] = None
+    with open(log_path) as f:
+        for line in f:
+            m = EVAL_RE.search(line)
+            if m:
+                try:
+                    last = float(m.group(1))
+                except ValueError:
+                    continue
+    return last
+
+
+def latest_ckpt_dir(out_dir: Path) -> Optional[Path]:
+    if not out_dir.exists():
+        return None
+    cks = sorted(
+        [p for p in out_dir.iterdir() if p.is_dir() and p.name.startswith("checkpoint-step-")],
+        key=lambda p: int(p.name.split("-")[-1]),
+    )
+    if cks:
+        return cks[-1]
+    if (out_dir / "adapter_model.safetensors").exists():
+        return out_dir
+    return None
+
+
+def best_grpo_ckpt(out_dir: Path) -> Optional[Path]:
+    if not out_dir.exists():
+        return None
+    cks = sorted(
+        [p for p in out_dir.iterdir() if p.is_dir() and p.name.startswith("checkpoint-")],
+        key=lambda p: int(p.name.split("-")[-1]) if p.name.split("-")[-1].isdigit() else -1,
+    )
+    if cks:
+        return cks[-1]
+    if (out_dir / "adapter_model.safetensors").exists():
+        return out_dir
+    return None
+
+
+# ---- subprocess wrappers ---------------------------------------------------
+
+
+def run_sft_phase(
+    *,
+    args: argparse.Namespace,
+    phase_dir: Path,
+    init_adapter: str,
+    num_cot_tokens: int,
+    max_steps: int,
+) -> Path:
+    """Launch one SFT phase at fixed k. Returns latest checkpoint path."""
+    phase_dir.mkdir(parents=True, exist_ok=True)
+    log_path = phase_dir / "train.log"
+    cmd = [
+        args.python_bin,
+        "-u",
+        str(SFT_SCRIPT),
+        "--model_name",
+        args.model_name,
+        "--train_jsonl",
+        str(TRAIN_JSONL),
+        "--eval_jsonl",
+        str(EVAL_JSONL),
+        "--output_dir",
+        str(phase_dir),
+        "--cache_dir",
+        args.cache_dir,
+        "--init_adapter_dir",
+        str(init_adapter),
+        "--seed",
+        str(args.seed),
+        "--gpu_id",
+        "0",
+        "--stage_i",
+        "3",
+        "--num_cot_tokens",
+        str(int(num_cot_tokens)),
+        "--latent_mode",
+        args.latent_mode,
+        "--total_empties_hint",
+        "20",
+        "--per_device_train_batch_size",
+        str(args.sft_bs),
+        "--gradient_accumulation_steps",
+        str(args.sft_ga),
+        "--num_epochs",
+        "256",
+        "--learning_rate",
+        str(args.sft_lr),
+        "--max_grad_norm",
+        "1.0",
+        "--logging_steps",
+        "25",
+        "--eval_steps",
+        "200",
+        "--save_steps",
+        "200",
+        "--eval_rows",
+        str(args.eval_rows),
+        "--max_completion_length",
+        "24",
+        "--limit_train_rows",
+        str(args.train_rows),
+        "--lora_r",
+        "32",
+        "--lora_alpha",
+        "64",
+        "--lora_dropout",
+        "0.05",
+        "--multi_value_oversample_factor",
+        str(args.sft_oversample),
+        "--max_steps",
+        str(int(max_steps)),
+    ]
+    if args.enable_gc:
+        cmd.append("--enable_gradient_checkpointing")
+
+    print(f"[adaptive-k] >>> SFT phase k={num_cot_tokens} max_steps={max_steps}", flush=True)
+    print(f"[adaptive-k]     init={init_adapter or '(BASE)'}", flush=True)
+    print(f"[adaptive-k]     out={phase_dir}", flush=True)
+    print(f"[adaptive-k]     log={log_path}", flush=True)
+
+    env = dict(os.environ)
+    env["CUDA_VISIBLE_DEVICES"] = str(args.gpu)
+    env["TOKENIZERS_PARALLELISM"] = "false"
+    env["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
+    env["HF_HOME"] = args.cache_dir
+    env["TRANSFORMERS_CACHE"] = args.cache_dir
+    with open(log_path, "w") as logf:
+        ret = subprocess.run(cmd, stdout=logf, stderr=subprocess.STDOUT, env=env)
+    if ret.returncode != 0:
+        raise RuntimeError(f"SFT phase k={num_cot_tokens} failed (exit {ret.returncode}); see {log_path}")
+
+    last = latest_ckpt_dir(phase_dir)
+    if last is None:
+        raise RuntimeError(f"No checkpoint produced under {phase_dir}")
+    return last
+
+
+def run_grpo_phase(
+    *,
+    args: argparse.Namespace,
+    phase_dir: Path,
+    init_adapter: str,
+    num_cot_tokens: int,
+    max_steps: int,
+) -> Optional[Path]:
+    phase_dir.mkdir(parents=True, exist_ok=True)
+    log_path = phase_dir / "train.log"
+    cmd = [
+        args.python_bin,
+        "-u",
+        str(GRPO_SCRIPT),
+        "--model_name",
+        args.model_name,
+        "--train_jsonl",
+        str(TRAIN_JSONL),
+        "--eval_jsonl",
+        str(EVAL_JSONL),
+        "--output_dir",
+        str(phase_dir),
+        "--cache_dir",
+        args.cache_dir,
+        "--init_adapter_dir",
+        str(init_adapter),
+        "--seed",
+        str(args.seed),
+        "--gpu_id",
+        "0",
+        "--stage_i",
+        "3",
+        "--num_cot_tokens",
+        str(int(num_cot_tokens)),
+        "--latent_mode",
+        args.latent_mode,
+        "--total_empties_hint",
+        "20",
+        "--per_device_train_batch_size",
+        str(args.grpo_bs),
+        "--gradient_accumulation_steps",
+        str(args.grpo_ga),
+        "--num_train_epochs",
+        "100",
+        "--learning_rate",
+        str(args.grpo_lr),
+        "--logging_steps",
+        "10",
+        "--save_steps",
+        "200",
+        "--eval_steps",
+        "150",
+        "--eval_rows",
+        str(args.eval_rows),
+        "--num_generations",
+        str(args.grpo_ng),
+        "--max_prompt_length",
+        str(args.grpo_max_prompt),
+        "--max_completion_length",
+        str(args.grpo_max_completion),
+        "--beta",
+        str(args.grpo_beta),
+        "--limit_train_rows",
+        str(args.train_rows),
+        "--lora_r",
+        "32",
+        "--lora_alpha",
+        "64",
+        "--lora_dropout",
+        "0.05",
+        "--max_steps",
+        str(int(max_steps)),
+    ]
+    if args.enable_gc:
+        cmd.append("--enable_gradient_checkpointing")
+    print(f"[adaptive-k] >>> GRPO phase k={num_cot_tokens} max_steps={max_steps}", flush=True)
+    print(f"[adaptive-k]     init={init_adapter}", flush=True)
+    print(f"[adaptive-k]     out={phase_dir}", flush=True)
+
+    env = dict(os.environ)
+    env["CUDA_VISIBLE_DEVICES"] = str(args.gpu)
+    env["TOKENIZERS_PARALLELISM"] = "false"
+    env["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
+    env["HF_HOME"] = args.cache_dir
+    env["TRANSFORMERS_CACHE"] = args.cache_dir
+    with open(log_path, "w") as logf:
+        ret = subprocess.run(cmd, stdout=logf, stderr=subprocess.STDOUT, env=env)
+    if ret.returncode != 0:
+        print(f"[adaptive-k] WARN: GRPO failed exit={ret.returncode}, see {log_path}", flush=True)
+    return best_grpo_ckpt(phase_dir)
+
+
+# ---- main loop -------------------------------------------------------------
+
+
+def main() -> None:
+    args = parse_args()
+    output_root = Path(args.output_root)
+    output_root.mkdir(parents=True, exist_ok=True)
+    state_path = output_root / "STATE.json"
+    pipeline_log = output_root / "PIPELINE.log"
+
+    def log(msg: str) -> None:
+        line = f"[{time.strftime('%H:%M:%S')}] {msg}"
+        print(line, flush=True)
+        with open(pipeline_log, "a") as f:
+            f.write(line + "\n")
+
+    log(f"===== ADAPTIVE-K {args.variant} on GPU {args.gpu} =====")
+    log(f"  start_k={args.start_k} max_k={args.max_k} steps_per_phase={args.steps_per_phase} max_phases_per_k={args.max_phases_per_k}")
+    log(f"  plateau_eps={args.plateau_eps} sft_lr={args.sft_lr} grpo_lr={args.grpo_lr}")
+    log(f"  output_root={output_root}")
+
+    history: List[dict] = []
+    cur_k = int(args.start_k)
+    cur_init: str = ""  # "" -> train from base
+    last_metric_at_k: Optional[float] = None
+    phases_at_k = 0
+    sft_phase_idx = 0
+
+    while cur_k <= int(args.max_k):
+        sft_phase_idx += 1
+        phase_dir = output_root / f"sft_phase{sft_phase_idx:02d}_k{cur_k}"
+        ckpt = run_sft_phase(
+            args=args,
+            phase_dir=phase_dir,
+            init_adapter=cur_init,
+            num_cot_tokens=cur_k,
+            max_steps=int(args.steps_per_phase),
+        )
+        metric = latest_eval_metric(phase_dir / "train.log")
+        log(
+            f"  phase{sft_phase_idx} k={cur_k} ckpt={ckpt.name}  eval_exact_set_match_rate={metric}"
+        )
+        history.append(
+            {
+                "phase": sft_phase_idx,
+                "k": cur_k,
+                "phase_dir": str(phase_dir),
+                "ckpt": str(ckpt),
+                "exact_set_match_rate": metric,
+            }
+        )
+        with open(state_path, "w") as f:
+            json.dump({"history": history, "cur_k": cur_k, "cur_ckpt": str(ckpt)}, f, indent=2)
+
+        cur_init = str(ckpt)
+        phases_at_k += 1
+
+        if cur_k >= int(args.max_k):
+            log(f"  reached max_k={args.max_k}, stopping SFT loop")
+            break
+
+        if last_metric_at_k is None or metric is None:
+            improvement = None
+        else:
+            improvement = float(metric) - float(last_metric_at_k)
+        log(f"  improvement_at_k={improvement}  phases_at_k={phases_at_k}/{args.max_phases_per_k}")
+
+        bump = False
+        if phases_at_k >= int(args.max_phases_per_k):
+            log("  hit max_phases_per_k, bumping k")
+            bump = True
+        elif improvement is not None and improvement < float(args.plateau_eps):
+            log(f"  improvement {improvement:.4f} < plateau_eps {args.plateau_eps:.4f}, bumping k")
+            bump = True
+
+        if bump:
+            cur_k += 1
+            last_metric_at_k = None
+            phases_at_k = 0
+        else:
+            last_metric_at_k = metric
+
+    log(f"===== final SFT k={cur_k} ckpt={cur_init} =====")
+    grpo_dir = output_root / f"grpo_k{cur_k}"
+    grpo_ckpt = run_grpo_phase(
+        args=args,
+        phase_dir=grpo_dir,
+        init_adapter=cur_init,
+        num_cot_tokens=cur_k,
+        max_steps=int(args.grpo_steps),
+    )
+    log(f"===== GRPO done ckpt={grpo_ckpt} =====")
+    with open(state_path, "w") as f:
+        json.dump(
+            {
+                "history": history,
+                "final_k": cur_k,
+                "final_sft_ckpt": cur_init,
+                "grpo_ckpt": str(grpo_ckpt) if grpo_ckpt else None,
+            },
+            f,
+            indent=2,
+        )
+    log(f"===== ADAPTIVE-K {args.variant} done =====")
+
+
+if __name__ == "__main__":
+    main()

_runs/adaptive_latent_baseline_sudoku_train.py

@@ -0,0 +1,534 @@
+#!/usr/bin/env python3
+"""Adaptive-k thought-token baseline (experiment D in the 2x2 ablation).
+
+Same single-stage, whole-puzzle setup as `simple_baseline_sudoku_train.py`
+(experiment C, the "strawman"). Same JSONL data, same chat template, same
+model, same LoRA. The ONLY difference is that this run inserts k recurrent
+thought tokens between the prompt and the next-token logits, and grows k
+on demand whenever the SFT loss plateaus.
+
+Algorithm:
+    k = 0   (start as the vanilla baseline)
+    repeat:
+        train SFT for `min_steps_per_k` steps with current k
+        if rolling_avg(loss[-w:]) - rolling_avg(loss[-2w:-w]) > -plateau_eps:
+            k += 1                          # grow capacity
+            if k > max_k: break
+        if loss has been steadily decreasing past `min_steps_per_k * 3`:
+            break                            # converged
+    save final adapter
+
+The recurrent_hidden mechanism is imported verbatim from
+`latent_multi_output_cell_policy.grpo_residual_projector_latent_train`
+(via `latent_batched_completion_ce_loss`). For k=0 the loss reduces to
+vanilla next-token CE, so the trajectory smoothly continues from the
+strawman.
+
+Reward / loss contract (see `simple_baseline_sudoku_train.py` for details):
+- supervision is token-level CE against the JSONL `completion` field
+  (the 20 ground-truth digits at the 20 empty cells, row-major).
+- this script is SFT-only; you can chain GRPO afterwards by passing the
+  saved adapter to `simple_baseline_sudoku_train.py --phase grpo`.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import math
+import os
+import sys
+import time
+from collections import deque
+from pathlib import Path
+from typing import Any, Dict, List, Tuple
+
+import torch
+import torch.nn.functional as F
+from peft import LoraConfig, PeftModel, get_peft_model
+from torch.optim import AdamW
+from transformers import AutoModelForCausalLM, AutoTokenizer, set_seed
+
+ROOT = Path(__file__).resolve().parent.parent
+if str(ROOT) not in sys.path:
+    sys.path.insert(0, str(ROOT))
+
+# Reuse helpers and the latent loss from the curriculum codebase. NO
+# re-implementation of the recurrent_hidden mechanism here.
+from multi_output_cell_policy.sft_multi_output_train import (  # type: ignore
+    load_jsonl_rows,
+    pick_dtype,
+)
+from latent_multi_output_cell_policy.sft_latent_multi_output_train import (  # type: ignore
+    latent_batched_completion_ce_loss,
+)
+from latent_multi_output_cell_policy.grpo_residual_projector_latent_train import (  # type: ignore
+    recurrent_hidden_next_token_logits_from_ids,
+)
+from _runs.simple_baseline_sudoku_train import (  # type: ignore
+    SYSTEM_PROMPT_STRAWMAN,
+    build_chat_prompt,
+    parse_int_list,
+)
+from multi_output_cell_policy.rewards import score_prediction_text  # type: ignore
+from multi_output_cell_policy.shared_multi_output_policy import (  # type: ignore
+    make_solved_grid_from_row,
+    stage_i_consistent_values,
+)
+from aligned_cell_policy.shared_cell_policy import build_cell_examples_from_row  # type: ignore
+
+
+# ---- Tokenization (mirror what latent_batched_completion_ce_loss expects) ---
+
+
+def tokenize_example(
+    tokenizer: Any,
+    raw_prompt: str,
+    completion_text: str,
+    max_prompt_length: int,
+    max_completion_length: int,
+) -> Dict[str, List[int]]:
+    prompt_text = build_chat_prompt(tokenizer, raw_prompt)
+    prompt_ids = tokenizer(
+        prompt_text, add_special_tokens=False
+    ).input_ids[-max_prompt_length:]
+    eos = tokenizer.eos_token or "<|endoftext|>"
+    completion_ids = tokenizer(
+        completion_text + eos, add_special_tokens=False
+    ).input_ids[:max_completion_length]
+    return {"prompt_ids": prompt_ids, "completion_ids": completion_ids}
+
+
+# ---- Eval (autoregressive greedy decode WITH k recurrent thought tokens) ---
+
+
+@torch.no_grad()
+def latent_greedy_generate(
+    model: torch.nn.Module,
+    tokenizer: Any,
+    prompt_text: str,
+    device: torch.device,
+    *,
+    num_cot_tokens: int,
+    max_new_tokens: int,
+) -> str:
+    enc = tokenizer(prompt_text, return_tensors="pt", add_special_tokens=False)
+    input_ids = enc["input_ids"].to(device)
+    attention_mask = enc["attention_mask"].to(device)
+    prompt_len = int(input_ids.shape[1])
+    eos_token_id = tokenizer.eos_token_id
+    for _ in range(int(max_new_tokens)):
+        logits = recurrent_hidden_next_token_logits_from_ids(
+            model, input_ids, attention_mask, int(max(0, num_cot_tokens))
+        )
+        next_id = int(torch.argmax(logits, dim=-1).item())
+        input_ids = torch.cat(
+            [input_ids, torch.tensor([[next_id]], device=device, dtype=input_ids.dtype)], dim=1
+        )
+        attention_mask = torch.cat(
+            [attention_mask, torch.ones((1, 1), device=device, dtype=attention_mask.dtype)], dim=1
+        )
+        if eos_token_id is not None and next_id == int(eos_token_id):
+            break
+    new_ids = input_ids[0, prompt_len:]
+    return tokenizer.decode(new_ids, skip_special_tokens=True).strip()
+
+
+@torch.no_grad()
+def run_eval(
+    model: torch.nn.Module,
+    tokenizer: Any,
+    eval_rows: List[Dict[str, Any]],
+    device: torch.device,
+    *,
+    num_cot_tokens: int,
+    max_new_tokens: int,
+    print_n: int = 3,
+    stage_i: int = 3,
+) -> Dict[str, float]:
+    """Apples-to-apples eval with the cell-policy framework (see strawman script).
+
+    The model emits the WHOLE puzzle (JSON list of integers) in one greedy
+    rollout with `num_cot_tokens` recurrent thought tokens prepended at each
+    step. We split that list into per-cell SINGLETON predictions and score
+    each cell with ``score_prediction_text`` against the i-consistent target
+    set at ``stage_i`` (default 3 — matches the S3 eval used for the v6
+    baseline and the latent champion).
+    """
+    model.eval()
+    total_cells = 0
+    parse_ok = 0.0
+    canonical_ok = 0.0
+    exact_set_match = 0.0
+    includes_gt = 0.0
+    precision_sum = 0.0
+    recall_sum = 0.0
+    cardinality_match_sum = 0.0
+    n_solve = 0
+    n_total_puzzles = 0
+    n_parse_fail_puzzles = 0
+    printed = 0
+    for row in eval_rows:
+        target_completion = parse_int_list(str(row["completion"]))
+        if target_completion is None:
+            continue
+        n_total_puzzles += 1
+        prompt_text = build_chat_prompt(tokenizer, str(row["prompt"]).strip())
+        gen = latent_greedy_generate(
+            model, tokenizer, prompt_text, device,
+            num_cot_tokens=num_cot_tokens, max_new_tokens=max_new_tokens,
+        )
+        pred_list = parse_int_list(gen)
+
+        try:
+            cells = build_cell_examples_from_row(row)
+            solved = make_solved_grid_from_row(row)
+        except Exception as e:
+            if printed < print_n:
+                print(f"[adaptive_k k={num_cot_tokens} eval] row skipped (no metadata): {e}", flush=True)
+                printed += 1
+            continue
+
+        row_all_exact = True
+        row_has_eval_cell = False
+        for idx, ex in enumerate(cells):
+            target_values = stage_i_consistent_values(
+                ex.grid, target_cell=ex.target_cell, stage_i=int(stage_i)
+            )
+            row_has_eval_cell = True
+            if pred_list is not None and idx < len(pred_list):
+                pred_text = json.dumps({"values": [int(pred_list[idx])]})
+            else:
+                pred_text = ""
+            info = score_prediction_text(
+                text=pred_text,
+                grid=ex.grid,
+                solved=solved,
+                target_cell=ex.target_cell,
+                stage_i=int(stage_i),
+                reward_good_value=1.0,
+                penalty_bad_value=1.75,
+                penalty_malformed=4.0,
+                penalty_empty=0.5,
+                penalty_singleton=1.5,
+            )
+            total_cells += 1
+            parse_ok += float(info["parse_ok"])
+            canonical_ok += float(info["strict_canonical"])
+            exact_set_match += float(info["exact_set_match"])
+            includes_gt += float(info["includes_ground_truth"])
+            precision_sum += float(info["value_precision"])
+            recall_sum += float(info["value_recall"])
+            if int(info["num_predicted_values"]) == int(len(target_values)):
+                cardinality_match_sum += 1.0
+            if float(info["exact_set_match"]) < 0.5:
+                row_all_exact = False
+        if row_has_eval_cell and row_all_exact:
+            n_solve += 1
+        if pred_list is None:
+            n_parse_fail_puzzles += 1
+        if printed < print_n:
+            head_pred = pred_list if pred_list is not None else "PARSE_FAIL"
+            print(
+                f"[adaptive_k k={num_cot_tokens} eval] target={target_completion}  pred={head_pred}  "
+                f"solve={int(row_all_exact and row_has_eval_cell)}  gen={gen!r}",
+                flush=True,
+            )
+            printed += 1
+    model.train()
+    return {
+        "n_total_cells": float(total_cells),
+        "n_total_puzzles": float(n_total_puzzles),
+        "parse_rate": float(parse_ok / max(1, total_cells)),
+        "strict_canonical_rate": float(canonical_ok / max(1, total_cells)),
+        "exact_set_match_rate": float(exact_set_match / max(1, total_cells)),
+        "includes_ground_truth_rate": float(includes_gt / max(1, total_cells)),
+        "value_precision": float(precision_sum / max(1, total_cells)),
+        "value_recall": float(recall_sum / max(1, total_cells)),
+        "cardinality_match_rate": float(cardinality_match_sum / max(1, total_cells)),
+        "puzzle_parse_fail_rate": float(n_parse_fail_puzzles / max(1, n_total_puzzles)),
+        "solve_rate": float(n_solve) / max(1, n_total_puzzles),
+    }
+
+
+# ---- Main loop --------------------------------------------------------------
+
+
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser()
+    p.add_argument("--model_name", type=str, default="Qwen/Qwen2.5-1.5B-Instruct")
+    p.add_argument("--train_jsonl", type=str, required=True)
+    p.add_argument("--eval_jsonl", type=str, required=True)
+    p.add_argument("--output_dir", type=str, required=True)
+    p.add_argument("--cache_dir", type=str, default=str(ROOT / ".hf_cache"))
+    p.add_argument("--init_adapter_dir", type=str, default="")
+    p.add_argument("--seed", type=int, default=0)
+
+    # Data
+    p.add_argument("--limit_train_rows", type=int, default=10000)
+    p.add_argument("--eval_rows", type=int, default=50)
+
+    # Train hyperparameters
+    p.add_argument("--per_device_train_batch_size", type=int, default=4)
+    p.add_argument("--gradient_accumulation_steps", type=int, default=2)
+    p.add_argument("--learning_rate", type=float, default=5e-5)
+    p.add_argument("--weight_decay", type=float, default=0.0)
+    p.add_argument("--max_steps", type=int, default=4000)
+    p.add_argument("--logging_steps", type=int, default=25)
+    p.add_argument("--save_steps", type=int, default=500)
+    p.add_argument("--eval_every_steps", type=int, default=500)
+    p.add_argument("--max_grad_norm", type=float, default=1.0)
+    p.add_argument("--max_completion_length", type=int, default=96)
+    p.add_argument("--max_prompt_length", type=int, default=1024)
+
+    # LoRA
+    p.add_argument("--lora_r", type=int, default=32)
+    p.add_argument("--lora_alpha", type=int, default=64)
+    p.add_argument("--lora_dropout", type=float, default=0.05)
+    p.add_argument("--enable_gradient_checkpointing", action="store_true")
+
+    # Adaptive-k schedule
+    p.add_argument("--start_k", type=int, default=0)
+    p.add_argument("--max_k", type=int, default=4)
+    p.add_argument(
+        "--min_steps_per_k",
+        type=int,
+        default=400,
+        help="Minimum SFT steps to spend at each k before considering an increment.",
+    )
+    p.add_argument(
+        "--plateau_window",
+        type=int,
+        default=100,
+        help="Sliding window (in steps) used to compute rolling-mean loss for plateau detection.",
+    )
+    p.add_argument(
+        "--plateau_eps",
+        type=float,
+        default=0.005,
+        help="If rolling_mean(loss[-w:]) - rolling_mean(loss[-2w:-w]) > -plateau_eps -> plateau detected.",
+    )
+    p.add_argument(
+        "--converged_eps",
+        type=float,
+        default=0.001,
+        help="If two consecutive plateau windows pass with delta within this band, we declare convergence and stop.",
+    )
+
+    return p.parse_args()
+
+
+def setup_model_and_tokenizer(args: argparse.Namespace, device: torch.device):
+    tokenizer = AutoTokenizer.from_pretrained(args.model_name, cache_dir=args.cache_dir, use_fast=True)
+    if tokenizer.pad_token_id is None:
+        tokenizer.pad_token = tokenizer.eos_token or "<|endoftext|>"
+    model = AutoModelForCausalLM.from_pretrained(
+        args.model_name, cache_dir=args.cache_dir,
+        torch_dtype=pick_dtype(), low_cpu_mem_usage=True,
+    )
+    if str(args.init_adapter_dir).strip():
+        model = PeftModel.from_pretrained(model, args.init_adapter_dir, is_trainable=True)
+    else:
+        lora = LoraConfig(
+            r=args.lora_r,
+            lora_alpha=args.lora_alpha,
+            lora_dropout=args.lora_dropout,
+            bias="none",
+            task_type="CAUSAL_LM",
+            target_modules=[
+                "q_proj", "k_proj", "v_proj", "o_proj",
+                "gate_proj", "up_proj", "down_proj",
+            ],
+        )
+        model = get_peft_model(model, lora)
+    if args.enable_gradient_checkpointing:
+        if hasattr(model, "gradient_checkpointing_enable"):
+            model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
+        if hasattr(model, "enable_input_require_grads"):
+            model.enable_input_require_grads()
+    if hasattr(model, "config"):
+        model.config.use_cache = False
+    model.to(device)
+    return model, tokenizer
+
+
+def detect_plateau(losses: deque, window: int, plateau_eps: float) -> Tuple[bool, float]:
+    if len(losses) < 2 * window:
+        return False, 0.0
+    arr = list(losses)
+    recent = arr[-window:]
+    prior = arr[-2 * window : -window]
+    delta = (sum(recent) / len(recent)) - (sum(prior) / len(prior))
+    # If delta > -plateau_eps, loss hasn't decreased fast enough -> plateau.
+    return (delta > -float(plateau_eps)), float(delta)
+
+
+def save_adapter(model: torch.nn.Module, tokenizer: Any, out: str) -> None:
+    os.makedirs(out, exist_ok=True)
+    if hasattr(model, "save_pretrained"):
+        model.save_pretrained(out)
+    if hasattr(tokenizer, "save_pretrained"):
+        tokenizer.save_pretrained(out)
+
+
+def main() -> None:
+    args = parse_args()
+    set_seed(int(args.seed))
+    os.makedirs(args.output_dir, exist_ok=True)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    train_rows = load_jsonl_rows(args.train_jsonl, limit_rows=int(args.limit_train_rows))
+    eval_rows = load_jsonl_rows(args.eval_jsonl, limit_rows=int(args.eval_rows))
+
+    print(f"[adaptive_k] loaded {len(train_rows)} train rows, {len(eval_rows)} eval rows", flush=True)
+
+    model, tokenizer = setup_model_and_tokenizer(args, device)
+    pad_id = int(tokenizer.pad_token_id)
+
+    # Pre-tokenize the train set once.
+    train_examples: List[Dict[str, Any]] = []
+    for row in train_rows:
+        try:
+            ex = tokenize_example(
+                tokenizer,
+                str(row["prompt"]).strip(),
+                str(row["completion"]).strip(),
+                int(args.max_prompt_length),
+                int(args.max_completion_length),
+            )
+            if ex["completion_ids"]:
+                train_examples.append(ex)
+        except Exception as e:  # noqa: BLE001
+            print(f"[adaptive_k] tokenize skip: {e}", flush=True)
+    print(f"[adaptive_k] tokenized {len(train_examples)} train examples", flush=True)
+
+    optimizer = AdamW(
+        (p for p in model.parameters() if p.requires_grad),
+        lr=float(args.learning_rate),
+        weight_decay=float(args.weight_decay),
+    )
+
+    bs = int(args.per_device_train_batch_size)
+    ga = int(args.gradient_accumulation_steps)
+    steps = 0
+    losses_per_step: List[float] = []
+    rolling: deque = deque(maxlen=2 * int(args.plateau_window) + 16)
+    k = int(args.start_k)
+    max_k = int(args.max_k)
+    steps_at_current_k = 0
+    grew_at: List[Tuple[int, int]] = []  # (step, new_k)
+
+    print(f"[adaptive_k] starting at k={k}", flush=True)
+    init_eval = run_eval(
+        model, tokenizer, eval_rows, device,
+        num_cot_tokens=k, max_new_tokens=int(args.max_completion_length),
+    )
+    print(f"[adaptive_k] init eval k={k}: {init_eval}", flush=True)
+
+    t0 = time.time()
+    rng_state = torch.Generator(device="cpu").manual_seed(int(args.seed))
+    perm = torch.randperm(len(train_examples), generator=rng_state).tolist()
+    cursor = 0
+
+    optimizer.zero_grad(set_to_none=True)
+    micro_in_step = 0
+    micro_loss_accum = 0.0
+
+    while steps < int(args.max_steps):
+        if cursor + bs > len(perm):
+            perm = torch.randperm(len(train_examples), generator=rng_state).tolist()
+            cursor = 0
+        batch_indices = perm[cursor : cursor + bs]
+        cursor += bs
+        batch = [train_examples[i] for i in batch_indices]
+
+        loss = latent_batched_completion_ce_loss(
+            model,
+            batch,
+            device,
+            num_cot_tokens=int(max(0, k)),
+            latent_mode="recurrent_hidden",
+            pad_token_id=pad_id,
+        ) / float(ga)
+        loss.backward()
+        micro_loss_accum += float(loss.detach().item()) * float(ga)
+        micro_in_step += 1
+
+        if micro_in_step >= ga:
+            torch.nn.utils.clip_grad_norm_(
+                (p for p in model.parameters() if p.requires_grad),
+                float(args.max_grad_norm),
+            )
+            optimizer.step()
+            optimizer.zero_grad(set_to_none=True)
+            steps += 1
+            steps_at_current_k += 1
+            avg_micro_loss = micro_loss_accum / float(ga)
+            losses_per_step.append(avg_micro_loss)
+            rolling.append(avg_micro_loss)
+            micro_in_step = 0
+            micro_loss_accum = 0.0
+
+            if steps % int(args.logging_steps) == 0:
+                w = int(args.plateau_window)
+                recent = list(rolling)[-w:] if len(rolling) >= w else list(rolling)
+                rec_mean = sum(recent) / max(1, len(recent))
+                elapsed = time.time() - t0
+                print(
+                    f"[adaptive_k] step={steps} k={k} loss={avg_micro_loss:.4f} "
+                    f"rolling_mean({len(recent)})={rec_mean:.4f} elapsed={elapsed:.0f}s",
+                    flush=True,
+                )
+
+            if steps % int(args.eval_every_steps) == 0:
+                ev = run_eval(
+                    model, tokenizer, eval_rows, device,
+                    num_cot_tokens=k, max_new_tokens=int(args.max_completion_length),
+                )
+                print(f"[adaptive_k] EVAL step={steps} k={k}: {ev}", flush=True)
+
+            if steps % int(args.save_steps) == 0:
+                save_adapter(model, tokenizer, os.path.join(args.output_dir, f"checkpoint-step-{steps:05d}"))
+
+            # Plateau check (only after `min_steps_per_k` at current k, and we
+            # have at least 2*plateau_window losses in the rolling buffer).
+            if steps_at_current_k >= int(args.min_steps_per_k):
+                plateau, delta = detect_plateau(rolling, int(args.plateau_window), float(args.plateau_eps))
+                if plateau and k < max_k:
+                    print(
+                        f"[adaptive_k] plateau detected at step={steps} k={k} delta={delta:+.4f} -> growing k -> {k+1}",
+                        flush=True,
+                    )
+                    k += 1
+                    steps_at_current_k = 0
+                    grew_at.append((steps, k))
+                    rolling.clear()  # restart plateau tracking after capacity bump
+                    save_adapter(model, tokenizer, os.path.join(args.output_dir, f"checkpoint-step-{steps:05d}-grow-k{k}"))
+                elif plateau and k >= max_k and abs(delta) < float(args.converged_eps):
+                    print(
+                        f"[adaptive_k] convergence at step={steps} k={k} delta={delta:+.4f} (max_k reached) -> stopping",
+                        flush=True,
+                    )
+                    break
+
+    final_dir = os.path.join(args.output_dir, "final")
+    save_adapter(model, tokenizer, final_dir)
+    final_eval = run_eval(
+        model, tokenizer, eval_rows, device,
+        num_cot_tokens=k, max_new_tokens=int(args.max_completion_length),
+    )
+    summary = {
+        "final_k": k,
+        "total_steps": steps,
+        "max_k": max_k,
+        "grew_at_steps": grew_at,
+        "final_eval": final_eval,
+        "training_seconds": time.time() - t0,
+    }
+    with open(os.path.join(args.output_dir, "summary.json"), "w", encoding="utf-8") as f:
+        json.dump(summary, f, indent=2)
+    print(f"[adaptive_k] DONE  summary={json.dumps(summary)}", flush=True)
+
+
+if __name__ == "__main__":
+    main()

_runs/add_variants_g_h.sh

@@ -0,0 +1,57 @@
+#!/usr/bin/env bash
+# Add 2 more variants on GPUs 6 and 7 to the active sweep.
+# Both seed from the lr5e5 lowsft S2 SFT step-3000 (the winning lineage at step 150).
+set -euo pipefail
+
+ROOT="${ROOT:-/home/ubuntu/curriculum_cot}"
+SWEEP_ROOT="${SWEEP_ROOT:-$(ls -dt ${ROOT}/_runs/baseline_1p5b_v4_*/ 2>/dev/null | head -1 | sed 's:/$::')}"
+PIPELINE="${ROOT}/_runs/baseline_1p5b_pipeline_v4.sh"
+
+[[ -d "${SWEEP_ROOT}" ]] || { echo "sweep root missing"; exit 1; }
+echo "Sweep: ${SWEEP_ROOT}"
+
+CKPT_LR5E5="${ROOT}/checkpoints/sudoku-9x9-20empty-baseline-1p5b-sweep/baseline_lr5e5_lowsft_v3/s2_sft_v3/checkpoint-step-03000"
+[[ -d "${CKPT_LR5E5}" ]] || { echo "missing init"; exit 1; }
+
+launch_variant() {
+  local gpu="$1" variant="$2" init="$3"
+  shift 3
+  local out="${SWEEP_ROOT}/${variant}"
+  mkdir -p "${out}"
+  local nohup_log="${out}/nohup.log"
+  printf 'GPU %s -> %s -> %s\n' "${gpu}" "${variant}" "${init}"
+  nohup env \
+    ROOT="${ROOT}" \
+    VARIANT="${variant}" \
+    GPU="${gpu}" \
+    S2_SFT_CKPT="${init}" \
+    OUTPUT_ROOT="${out}" \
+    USE_WANDB=0 \
+    WANDB_MODE=offline \
+    "$@" \
+    bash "${PIPELINE}" \
+    </dev/null >"${nohup_log}" 2>&1 &
+  local pid=$!
+  printf '  pid=%s   log=%s\n' "${pid}" "${nohup_log}"
+  echo "${pid} ${gpu} ${variant}" >> "${SWEEP_ROOT}/PIDS.txt"
+  disown "${pid}" 2>/dev/null || true
+}
+
+# pipe_g: lr5e5 lineage, faster GRPO LR (1e-5) to push convergence
+launch_variant 6 pipe_g_lr5e5_grpo1e5        "${CKPT_LR5E5}" GRPO_LR=1e-5 SFT_LR_S3=2e-5 PENALTY_SINGLETON=1.5
+
+# pipe_h: lr5e5 lineage, lower singleton penalty (1.0) to test if 1.5 hurts
+launch_variant 7 pipe_h_lr5e5_grpo5e6_sngl10 "${CKPT_LR5E5}" GRPO_LR=5e-6 SFT_LR_S3=2e-5 PENALTY_SINGLETON=1.0
+
+# Update sweep README
+cat >>"${SWEEP_ROOT}/SWEEP_README.md" <<EOF
+
+## Added at $(date '+%H:%M:%S')
+
+| GPU | variant | S2 init | GRPO LR | S3 SFT LR | penalty_singleton |
+| ---: | --- | --- | ---: | ---: | ---: |
+| 6 | pipe_g_lr5e5_grpo1e5         | lr5e5_lowsft step-3000 | 1e-5 | 2e-5 | 1.5 |
+| 7 | pipe_h_lr5e5_grpo5e6_sngl10  | lr5e5_lowsft step-3000 | 5e-6 | 2e-5 | 1.0 |
+EOF
+
+echo "Done. Now running 8 variants on GPUs 0..7."

_runs/add_variants_i_j_k_l.sh

@@ -0,0 +1,94 @@
+#!/usr/bin/env bash
+# Add 4 high-throughput variants on freed GPUs 0,2,3,4.
+# 3 variants fast-forward to S3 SFT (since S2 GRPO is plateau-stuck on baseline).
+# 1 variant tries an aggressive 10x GRPO LR to break the S2 plateau.
+set -euo pipefail
+
+ROOT="${ROOT:-/home/ubuntu/curriculum_cot}"
+SWEEP_ROOT="${SWEEP_ROOT:-$(ls -dt ${ROOT}/_runs/baseline_1p5b_v4_*/ 2>/dev/null | head -1 | sed 's:/$::')}"
+PIPELINE="${ROOT}/_runs/baseline_1p5b_pipeline_v4.sh"
+
+[[ -d "${SWEEP_ROOT}" ]] || { echo "sweep root missing"; exit 1; }
+echo "Sweep: ${SWEEP_ROOT}"
+
+CKPT_LR1E4="${ROOT}/checkpoints/sudoku-9x9-20empty-baseline-1p5b-sweep/baseline_lr1e4_lowsft_v3/s2_sft_v3/checkpoint-step-03000"
+CKPT_LR5E5="${ROOT}/checkpoints/sudoku-9x9-20empty-baseline-1p5b-sweep/baseline_lr5e5_lowsft_v3/s2_sft_v3/checkpoint-step-03000"
+[[ -d "${CKPT_LR5E5}" ]] || { echo "missing init lr5e5"; exit 1; }
+[[ -d "${CKPT_LR1E4}" ]] || { echo "missing init lr1e4"; exit 1; }
+
+launch_variant() {
+  local gpu="$1" variant="$2" init="$3"
+  shift 3
+  local out="${SWEEP_ROOT}/${variant}"
+  mkdir -p "${out}"
+  local nohup_log="${out}/nohup.log"
+  printf 'GPU %s -> %s -> %s\n' "${gpu}" "${variant}" "${init}"
+  nohup env \
+    ROOT="${ROOT}" \
+    VARIANT="${variant}" \
+    GPU="${gpu}" \
+    S2_SFT_CKPT="${init}" \
+    OUTPUT_ROOT="${out}" \
+    USE_WANDB=0 \
+    WANDB_MODE=offline \
+    "$@" \
+    bash "${PIPELINE}" \
+    </dev/null >"${nohup_log}" 2>&1 &
+  local pid=$!
+  printf '  pid=%s   log=%s\n' "${pid}" "${nohup_log}"
+  echo "${pid} ${gpu} ${variant}" >> "${SWEEP_ROOT}/PIDS.txt"
+  disown "${pid}" 2>/dev/null || true
+}
+
+# pipe_i (GPU 0): fast-forward to S3 SFT from lr5e5 lowsft step-3000.
+#   high-throughput: no GC, bs=32x1, larger eval batches.
+launch_variant 0 pipe_i_s3sft_lr5e5_fast "${CKPT_LR5E5}" \
+  START_PHASE=s3_sft S3_SFT_INIT="${CKPT_LR5E5}" \
+  SFT_LR_S3=2e-5 SFT_BS=32 SFT_GA=1 \
+  GRPO_LR=5e-6 GRPO_BS=32 GRPO_GA=1 GRPO_NG=8 \
+  USE_GC=0
+
+# pipe_j (GPU 2): fast-forward to S3 SFT from lr5e5 with lower LR for stability.
+launch_variant 2 pipe_j_s3sft_lr5e5_lr1e5 "${CKPT_LR5E5}" \
+  START_PHASE=s3_sft S3_SFT_INIT="${CKPT_LR5E5}" \
+  SFT_LR_S3=1e-5 SFT_BS=32 SFT_GA=1 \
+  GRPO_LR=5e-6 GRPO_BS=32 GRPO_GA=1 GRPO_NG=8 \
+  USE_GC=0
+
+# pipe_k (GPU 3): fast-forward to S3 SFT from lr1e4 lineage (mirror of i but other init).
+launch_variant 3 pipe_k_s3sft_lr1e4_fast "${CKPT_LR1E4}" \
+  START_PHASE=s3_sft S3_SFT_INIT="${CKPT_LR1E4}" \
+  SFT_LR_S3=2e-5 SFT_BS=32 SFT_GA=1 \
+  GRPO_LR=5e-6 GRPO_BS=32 GRPO_GA=1 GRPO_NG=8 \
+  USE_GC=0
+
+# pipe_l (GPU 4): aggressive 10x GRPO LR + 16 generations, push past S2 plateau.
+launch_variant 4 pipe_l_lr5e5_grpo5e5_ng16 "${CKPT_LR5E5}" \
+  START_PHASE=s2_grpo \
+  GRPO_LR=5e-5 GRPO_BS=16 GRPO_GA=1 GRPO_NG=16 \
+  PENALTY_SINGLETON=1.5 \
+  SFT_LR_S3=2e-5 SFT_BS=32 SFT_GA=1 \
+  USE_GC=0
+
+cat >>"${SWEEP_ROOT}/SWEEP_README.md" <<EOF
+
+## Added at $(date '+%H:%M:%S') — high-throughput / S3 fast-forward
+
+S2 GRPO plateaued at solve=0.14 (lr5e5 lineage) or 0.05 (lr1e4 lineage) for all
+of pipe_a/b/c/d/e — bit-identical evals from step 150 to 450. The per-cell
+exact ceiling (~0.91) caps puzzle solve at ~0.91^20 ~= 0.14 regardless of
+GRPO. Real lever is S3 SFT on harder cells (multi-value).
+
+Killed pipe_a, pipe_c, pipe_d, pipe_e (flat). Launched 4 replacements with
+USE_GC=0 (gradient checkpointing OFF — we have 80 GB headroom) and bs=32x1
+for ~2-3x throughput per GPU.
+
+| GPU | variant | start phase | init | SFT LR (S3) | GRPO LR | bs | ng |
+| ---: | --- | --- | --- | ---: | ---: | ---: | ---: |
+| 0 | pipe_i_s3sft_lr5e5_fast       | s3_sft  | lr5e5 step-3000 | 2e-5 | 5e-6 | 32 | 8 |
+| 2 | pipe_j_s3sft_lr5e5_lr1e5      | s3_sft  | lr5e5 step-3000 | 1e-5 | 5e-6 | 32 | 8 |
+| 3 | pipe_k_s3sft_lr1e4_fast       | s3_sft  | lr1e4 step-3000 | 2e-5 | 5e-6 | 32 | 8 |
+| 4 | pipe_l_lr5e5_grpo5e5_ng16     | s2_grpo | lr5e5 step-3000 |   -  | 5e-5 | 16 | 16 |
+EOF
+
+echo "Done. Now running 8 variants on GPUs 0..7."

_runs/baseline_1p5b_pipeline_v4.sh

@@ -0,0 +1,328 @@
+#!/usr/bin/env bash
+# 1.5B vanilla baseline: S2 GRPO -> S3 SFT -> S3 GRPO, single GPU.
+# Optionally pre-pends extra S2 SFT steps if EXTRA_S2_SFT_STEPS>0.
+#
+# Required env:
+#   VARIANT          variant name (used in dirs / wandb)
+#   GPU              CUDA index for this variant (0..7)
+#   S2_SFT_CKPT      path to S2 SFT LoRA adapter (uses this as S2 GRPO init)
+#
+# Optional env:
+#   ROOT             default /home/ubuntu/curriculum_cot
+#   PYTHON_BIN       default /opt/pytorch/bin/python
+#   OUTPUT_ROOT      default $ROOT/_runs/baseline_1p5b_v4_$(date)/$VARIANT
+#   MODEL_NAME       default Qwen/Qwen2.5-1.5B-Instruct
+#   GRPO_LR          default 5e-6
+#   GRPO_BETA        default 0.0
+#   GRPO_NG          default 8
+#   GRPO_BS          default 16
+#   GRPO_GA          default 2
+#   GRPO_PROMPT      default 768
+#   GRPO_COMPL       default 24
+#   PENALTY_SINGLETON default 1.5
+#   PENALTY_BAD      default 1.0
+#   REWARD_GOOD      default 1.25
+#   PENALTY_MAL      default 4.0
+#   PENALTY_EMPTY    default 0.5
+#   SFT_LR_S3        default 2e-5
+#   SFT_BS           default 16
+#   SFT_GA           default 2
+#   VALUE_TARGET     default 0.98
+#   S2_GRPO_MAX_STEPS  default 1200  (pipeline budget)
+#   S3_SFT_MAX_STEPS   default 2400
+#   S3_GRPO_MAX_STEPS  default 1500
+#   EXTRA_S2_SFT_STEPS default 0    (extra S2 SFT steps before S2 GRPO)
+#   EXTRA_S2_SFT_LR    default 1e-5
+#   EVAL_ROWS          default 100
+#   TRAIN_ROWS         default 10000
+#   USE_WANDB          default 0
+#   WANDB_PROJECT      default sudoku-baseline-1p5b-v4
+#   WANDB_MODE         default offline
+#   PHASE_WALL_SECS    default 0    (no phase wallclock cap)
+#   START_PHASE        default s2_grpo  (one of: s2_sft_extra,s2_grpo,s3_sft,s3_grpo)
+#   S3_SFT_INIT        if START_PHASE=s3_sft, S3-SFT init adapter (overrides S2 GRPO output)
+#   S3_GRPO_INIT       if START_PHASE=s3_grpo, S3-GRPO init adapter
+#   USE_GC             default 0  (1 to enable gradient checkpointing; we usually have memory)
+
+set -euo pipefail
+
+ROOT="${ROOT:-/home/ubuntu/curriculum_cot}"
+PYTHON_BIN="${PYTHON_BIN:-/opt/pytorch/bin/python}"
+SFT_SCRIPT="${ROOT}/multi_output_cell_policy/sft_multi_output_train.py"
+GRPO_SCRIPT="${ROOT}/multi_output_cell_policy/grpo_multi_output_train.py"
+
+: "${VARIANT:?VARIANT required}"
+: "${GPU:?GPU required}"
+: "${S2_SFT_CKPT:?S2_SFT_CKPT required}"
+
+OUTPUT_ROOT="${OUTPUT_ROOT:-${ROOT}/_runs/baseline_1p5b_v4_$(date +%Y%m%d_%H%M%S)/${VARIANT}}"
+MODEL_NAME="${MODEL_NAME:-Qwen/Qwen2.5-1.5B-Instruct}"
+
+GRPO_LR="${GRPO_LR:-5e-6}"
+GRPO_BETA="${GRPO_BETA:-0.0}"
+GRPO_NG="${GRPO_NG:-8}"
+GRPO_BS="${GRPO_BS:-16}"
+GRPO_GA="${GRPO_GA:-2}"
+GRPO_PROMPT="${GRPO_PROMPT:-768}"
+GRPO_COMPL="${GRPO_COMPL:-24}"
+PENALTY_SINGLETON="${PENALTY_SINGLETON:-1.5}"
+PENALTY_BAD="${PENALTY_BAD:-1.0}"
+PENALTY_MAL="${PENALTY_MAL:-4.0}"
+PENALTY_EMPTY="${PENALTY_EMPTY:-0.5}"
+REWARD_GOOD="${REWARD_GOOD:-1.25}"
+PENALTY_MISSING="${PENALTY_MISSING:-0.0}"
+EXACT_MATCH_BONUS="${EXACT_MATCH_BONUS:-0.0}"
+CARD_MISMATCH_PEN="${CARD_MISMATCH_PEN:-0.0}"
+SFT_OVERSAMPLE="${SFT_OVERSAMPLE:-1}"
+SFT_TGT_MIN="${SFT_TGT_MIN:-0}"
+SFT_TGT_MAX="${SFT_TGT_MAX:-0}"
+SFT_LR_S3="${SFT_LR_S3:-2e-5}"
+SFT_BS="${SFT_BS:-16}"
+SFT_GA="${SFT_GA:-2}"
+VALUE_TARGET="${VALUE_TARGET:-0.98}"
+S2_GRPO_MAX_STEPS="${S2_GRPO_MAX_STEPS:-1200}"
+S3_SFT_MAX_STEPS="${S3_SFT_MAX_STEPS:-2400}"
+S3_GRPO_MAX_STEPS="${S3_GRPO_MAX_STEPS:-1500}"
+EXTRA_S2_SFT_STEPS="${EXTRA_S2_SFT_STEPS:-0}"
+EXTRA_S2_SFT_LR="${EXTRA_S2_SFT_LR:-1e-5}"
+EVAL_ROWS="${EVAL_ROWS:-100}"
+TRAIN_ROWS="${TRAIN_ROWS:-10000}"
+USE_WANDB="${USE_WANDB:-0}"
+WANDB_PROJECT="${WANDB_PROJECT:-sudoku-baseline-1p5b-v4}"
+WANDB_MODE="${WANDB_MODE:-offline}"
+PHASE_WALL_SECS="${PHASE_WALL_SECS:-0}"
+START_PHASE="${START_PHASE:-s2_grpo}"
+S3_SFT_INIT="${S3_SFT_INIT:-}"
+S3_GRPO_INIT="${S3_GRPO_INIT:-}"
+USE_GC="${USE_GC:-0}"
+
+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"
+
+mkdir -p "${OUTPUT_ROOT}"
+PIPELINE_LOG="${OUTPUT_ROOT}/PIPELINE.log"
+
+ts() { date +'%H:%M:%S'; }
+log() { printf '[%s] %s\n' "$(ts)" "$*" | tee -a "${PIPELINE_LOG}" >&2; }
+
+latest_ckpt_step() {
+  local d="$1"
+  shopt -s nullglob
+  local cks=("${d}"/checkpoint-step-*)
+  shopt -u nullglob
+  (( ${#cks[@]} == 0 )) && return 1
+  printf '%s\n' "${cks[@]}" | sort -V | tail -n 1
+}
+
+best_grpo_adapter() {
+  local d="$1"
+  if [[ -f "${d}/adapter_model.safetensors" ]]; then
+    printf '%s\n' "${d}"; return 0
+  fi
+  local best="" step=-1
+  shopt -s nullglob
+  for c in "${d}"/checkpoint-*; do
+    [[ -d "$c" ]] || continue
+    [[ -f "$c/adapter_model.safetensors" ]] || continue
+    local n="${c##*checkpoint-}"
+    if [[ "$n" =~ ^[0-9]+$ ]] && (( 10#${n} >= step )); then
+      step=$((10#${n})); best="$c"
+    fi
+  done
+  shopt -u nullglob
+  [[ -n "$best" ]] || return 1
+  printf '%s\n' "$best"
+}
+
+if [[ ! -f "${TRAIN_JSONL}" || ! -f "${EVAL_JSONL}" ]]; then
+  log "ERROR: missing dataset jsonls (${TRAIN_JSONL} / ${EVAL_JSONL})."
+  exit 1
+fi
+
+export CUDA_VISIBLE_DEVICES="${GPU}"
+export TOKENIZERS_PARALLELISM=false
+export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
+export HF_HOME="${ROOT}/.hf_cache"
+export TRANSFORMERS_CACHE="${ROOT}/.hf_cache"
+export WANDB_MODE="${WANDB_MODE}"
+
+run_sft() {
+  local stage="$1" init_adapter="$2" out_dir="$3" lr="$4" max_steps="$5" tag="$6"
+  mkdir -p "${out_dir}"
+  log "=== Stage ${stage} SFT (${tag}) lr=${lr} max_steps=${max_steps} bs=${SFT_BS}x${SFT_GA} GC=${USE_GC} init=${init_adapter} ==="
+  log "    out=${out_dir}"
+  local extra=()
+  if [[ "${USE_WANDB}" == "1" ]]; then
+    extra+=(--use_wandb --wandb_project "${WANDB_PROJECT}" \
+            --wandb_run_name "${VARIANT}_${tag}" --wandb_mode "${WANDB_MODE}")
+  fi
+  if [[ "${USE_GC}" == "1" ]]; then
+    extra+=(--enable_gradient_checkpointing)
+  fi
+  "${PYTHON_BIN}" -u "${SFT_SCRIPT}" \
+    --model_name "${MODEL_NAME}" \
+    --train_jsonl "${TRAIN_JSONL}" \
+    --eval_jsonl "${EVAL_JSONL}" \
+    --output_dir "${out_dir}" \
+    --cache_dir "${ROOT}/.hf_cache" \
+    --init_adapter_dir "${init_adapter}" \
+    --seed 0 \
+    --gpu_id 0 \
+    --stage_i "${stage}" \
+    --total_empties_hint 20 \
+    --per_device_train_batch_size "${SFT_BS}" \
+    --gradient_accumulation_steps "${SFT_GA}" \
+    --num_epochs 256 \
+    --learning_rate "${lr}" \
+    --max_grad_norm 1.0 \
+    --logging_steps 25 \
+    --eval_steps 150 \
+    --save_steps 200 \
+    --eval_rows "${EVAL_ROWS}" \
+    --max_completion_length 24 \
+    --limit_train_rows "${TRAIN_ROWS}" \
+    --lora_r 32 --lora_alpha 64 --lora_dropout 0.05 \
+    --eval_value_precision_stop "${VALUE_TARGET}" \
+    --eval_value_recall_stop "${VALUE_TARGET}" \
+    --eval_exact_set_match_stop 0 \
+    --eval_solve_rate_stop 0 \
+    --min_steps_before_stop 100 \
+    --max_wall_clock_seconds "${PHASE_WALL_SECS}" \
+    --max_steps "${max_steps}" \
+    --multi_value_oversample_factor "${SFT_OVERSAMPLE}" \
+    --train_target_size_min "${SFT_TGT_MIN}" \
+    --train_target_size_max "${SFT_TGT_MAX}" \
+    "${extra[@]}" 2>&1 | tee "${out_dir}/train.log"
+}
+
+run_grpo() {
+  local stage="$1" init_adapter="$2" out_dir="$3" max_steps="$4" tag="$5"
+  mkdir -p "${out_dir}"
+  log "=== Stage ${stage} GRPO (${tag}) lr=${GRPO_LR} ng=${GRPO_NG} bs=${GRPO_BS}x${GRPO_GA} prompt=${GRPO_PROMPT} GC=${USE_GC} max_steps=${max_steps} init=${init_adapter} ==="
+  log "    rewards: good=${REWARD_GOOD} bad=${PENALTY_BAD} mal=${PENALTY_MAL} empty=${PENALTY_EMPTY} sngl=${PENALTY_SINGLETON} missing=${PENALTY_MISSING} exact_b=${EXACT_MATCH_BONUS} card_pen=${CARD_MISMATCH_PEN}"
+  log "    out=${out_dir}"
+  local extra=()
+  if [[ "${USE_WANDB}" == "1" ]]; then
+    extra+=(--use_wandb --wandb_project "${WANDB_PROJECT}" \
+            --wandb_run_name "${VARIANT}_${tag}" --wandb_mode "${WANDB_MODE}")
+  fi
+  if [[ "${USE_GC}" == "1" ]]; then
+    extra+=(--enable_gradient_checkpointing)
+  fi
+  "${PYTHON_BIN}" -u "${GRPO_SCRIPT}" \
+    --model_name "${MODEL_NAME}" \
+    --train_jsonl "${TRAIN_JSONL}" \
+    --eval_jsonl "${EVAL_JSONL}" \
+    --output_dir "${out_dir}" \
+    --cache_dir "${ROOT}/.hf_cache" \
+    --init_adapter_dir "${init_adapter}" \
+    --seed 0 \
+    --gpu_id 0 \
+    --stage_i "${stage}" \
+    --total_empties_hint 20 \
+    --per_device_train_batch_size "${GRPO_BS}" \
+    --gradient_accumulation_steps "${GRPO_GA}" \
+    --num_train_epochs 100 \
+    --learning_rate "${GRPO_LR}" \
+    --logging_steps 10 \
+    --save_steps 200 \
+    --eval_steps 150 \
+    --eval_rows "${EVAL_ROWS}" \
+    --num_generations "${GRPO_NG}" \
+    --max_prompt_length "${GRPO_PROMPT}" \
+    --max_completion_length "${GRPO_COMPL}" \
+    --beta "${GRPO_BETA}" \
+    --limit_train_rows "${TRAIN_ROWS}" \
+    --lora_r 32 --lora_alpha 64 --lora_dropout 0.05 \
+    --reward_good_value "${REWARD_GOOD}" \
+    --penalty_bad_value "${PENALTY_BAD}" \
+    --penalty_malformed "${PENALTY_MAL}" \
+    --penalty_empty "${PENALTY_EMPTY}" \
+    --penalty_singleton "${PENALTY_SINGLETON}" \
+    --penalty_missing "${PENALTY_MISSING}" \
+    --exact_match_bonus "${EXACT_MATCH_BONUS}" \
+    --cardinality_mismatch_penalty "${CARD_MISMATCH_PEN}" \
+    --eval_value_precision_stop "${VALUE_TARGET}" \
+    --eval_value_recall_stop "${VALUE_TARGET}" \
+    --eval_solve_rate_stop 0 \
+    --min_steps_before_stop 100 \
+    --max_wall_clock_seconds "${PHASE_WALL_SECS}" \
+    --max_steps "${max_steps}" \
+    "${extra[@]}" 2>&1 | tee "${out_dir}/train.log"
+}
+
+log "===== ${VARIANT} on GPU ${GPU} ====="
+log "S2 SFT init: ${S2_SFT_CKPT}"
+log "START_PHASE=${START_PHASE}  GRPO_LR=${GRPO_LR}  SFT_LR_S3=${SFT_LR_S3}  PENALTY_SINGLETON=${PENALTY_SINGLETON}  USE_GC=${USE_GC}"
+log "  EXTRA_S2_SFT_STEPS=${EXTRA_S2_SFT_STEPS}  GRPO_BS=${GRPO_BS}x${GRPO_GA}  SFT_BS=${SFT_BS}x${SFT_GA}  GRPO_NG=${GRPO_NG}"
+
+S2_SFT_DIR_FOR_GRPO="${S2_SFT_CKPT}"
+S2_GRPO_ADAPTER=""
+S3_SFT_INIT_RESOLVED=""
+S3_GRPO_INIT_RESOLVED=""
+
+phase_idx() {
+  case "$1" in
+    s2_sft_extra) echo 1 ;;
+    s2_grpo) echo 2 ;;
+    s3_sft) echo 3 ;;
+    s3_grpo) echo 4 ;;
+    *) echo 2 ;;
+  esac
+}
+START_IDX="$(phase_idx "${START_PHASE}")"
+
+if (( START_IDX <= 1 )) && (( EXTRA_S2_SFT_STEPS > 0 )); then
+  S2_SFT_EXTRA_DIR="${OUTPUT_ROOT}/s2_sft_extra"
+  run_sft 2 "${S2_SFT_CKPT}" "${S2_SFT_EXTRA_DIR}" "${EXTRA_S2_SFT_LR}" "${EXTRA_S2_SFT_STEPS}" "s2sft_extra"
+  if NEW_CKPT="$(latest_ckpt_step "${S2_SFT_EXTRA_DIR}")"; then
+    log ">>> Extra S2 SFT ckpt: ${NEW_CKPT}"
+    S2_SFT_DIR_FOR_GRPO="${NEW_CKPT}"
+  else
+    log "WARN: no new S2 SFT ckpt produced; falling back to ${S2_SFT_CKPT}"
+  fi
+fi
+
+if (( START_IDX <= 2 )); then
+  S2_GRPO_DIR="${OUTPUT_ROOT}/s2_grpo"
+  run_grpo 2 "${S2_SFT_DIR_FOR_GRPO}" "${S2_GRPO_DIR}" "${S2_GRPO_MAX_STEPS}" "s2grpo"
+  S2_GRPO_ADAPTER="$(best_grpo_adapter "${S2_GRPO_DIR}")"
+  if [[ -z "${S2_GRPO_ADAPTER}" ]]; then
+    log "ERROR: no S2 GRPO adapter under ${S2_GRPO_DIR}"; exit 1
+  fi
+  log ">>> S2 GRPO adapter: ${S2_GRPO_ADAPTER}"
+  S3_SFT_INIT_RESOLVED="${S2_GRPO_ADAPTER}"
+elif (( START_IDX == 3 )); then
+  if [[ -z "${S3_SFT_INIT}" ]]; then
+    log "ERROR: START_PHASE=s3_sft but S3_SFT_INIT is empty"; exit 1
+  fi
+  S3_SFT_INIT_RESOLVED="${S3_SFT_INIT}"
+  log ">>> Skipping to S3 SFT, init=${S3_SFT_INIT_RESOLVED}"
+fi
+
+if (( START_IDX <= 3 )); then
+  S3_SFT_DIR="${OUTPUT_ROOT}/s3_sft"
+  run_sft 3 "${S3_SFT_INIT_RESOLVED}" "${S3_SFT_DIR}" "${SFT_LR_S3}" "${S3_SFT_MAX_STEPS}" "s3sft"
+  S3_SFT_CKPT="$(latest_ckpt_step "${S3_SFT_DIR}")"
+  if [[ -z "${S3_SFT_CKPT}" ]]; then
+    log "ERROR: no S3 SFT ckpt under ${S3_SFT_DIR}"; exit 1
+  fi
+  log ">>> S3 SFT ckpt: ${S3_SFT_CKPT}"
+  S3_GRPO_INIT_RESOLVED="${S3_SFT_CKPT}"
+elif (( START_IDX == 4 )); then
+  if [[ -z "${S3_GRPO_INIT}" ]]; then
+    log "ERROR: START_PHASE=s3_grpo but S3_GRPO_INIT is empty"; exit 1
+  fi
+  S3_GRPO_INIT_RESOLVED="${S3_GRPO_INIT}"
+  log ">>> Skipping to S3 GRPO, init=${S3_GRPO_INIT_RESOLVED}"
+fi
+
+S3_GRPO_DIR="${OUTPUT_ROOT}/s3_grpo"
+run_grpo 3 "${S3_GRPO_INIT_RESOLVED}" "${S3_GRPO_DIR}" "${S3_GRPO_MAX_STEPS}" "s3grpo"
+S3_GRPO_ADAPTER="$(best_grpo_adapter "${S3_GRPO_DIR}")"
+if [[ -z "${S3_GRPO_ADAPTER}" ]]; then
+  log "ERROR: no S3 GRPO adapter under ${S3_GRPO_DIR}"; exit 1
+fi
+log ">>> S3 GRPO adapter: ${S3_GRPO_ADAPTER}"
+
+log "===== ${VARIANT} DONE — final S3 GRPO adapter at ${S3_GRPO_ADAPTER} ====="

_runs/eval_strawman_cellpolicy.py

@@ -0,0 +1,132 @@
+#!/usr/bin/env python3
+"""Re-evaluate any strawman / adaptive-k checkpoint using the cell-policy metric.
+
+This is a thin CLI wrapper that:
+
+  1. Loads a base model + LoRA adapter.
+  2. Runs the same scoring procedure as
+     ``multi_output_cell_policy/sft_multi_output_train.py::run_eval``,
+     i.e. for each puzzle it uses ``build_cell_examples_from_row`` to iterate
+     over empty cells in row-major order and scores each predicted value
+     with ``score_prediction_text`` against the i-consistent target set at
+     ``--stage_i`` (default 3, matching the S3 eval reported in the rebuttal).
+  3. The only difference vs the cell-policy is that the model emits the whole
+     puzzle in ONE forward pass, then the predicted list is split into
+     per-cell singletons.
+
+Use ``--kind strawman`` for vanilla LoRA models (``simple_baseline_sudoku_train.py``)
+and ``--kind adaptive_k --num_cot_tokens K`` for recurrent-hidden adaptive-k
+models (``adaptive_latent_baseline_sudoku_train.py``).
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from pathlib import Path
+from typing import Any, Dict, List
+
+import torch
+from peft import PeftModel
+from transformers import AutoModelForCausalLM, AutoTokenizer, set_seed
+
+ROOT = Path(__file__).resolve().parent.parent
+if str(ROOT) not in sys.path:
+    sys.path.insert(0, str(ROOT))
+
+from multi_output_cell_policy.sft_multi_output_train import (  # type: ignore  # noqa: E402
+    load_jsonl_rows,
+    pick_dtype,
+)
+from _runs.simple_baseline_sudoku_train import (  # type: ignore  # noqa: E402
+    run_eval as run_eval_strawman,
+)
+from _runs.adaptive_latent_baseline_sudoku_train import (  # type: ignore  # noqa: E402
+    run_eval as run_eval_adaptive_k,
+)
+
+
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser()
+    p.add_argument("--kind", choices=["strawman", "adaptive_k"], required=True)
+    p.add_argument("--model_name", default="Qwen/Qwen2.5-1.5B-Instruct")
+    p.add_argument("--adapter_dir", required=True)
+    p.add_argument("--eval_jsonl", required=True)
+    p.add_argument("--cache_dir", default=str(ROOT / ".hf_cache"))
+    p.add_argument("--eval_rows", type=int, default=100)
+    p.add_argument("--max_completion_length", type=int, default=96)
+    p.add_argument("--stage_i", type=int, default=3)
+    p.add_argument(
+        "--num_cot_tokens",
+        type=int,
+        default=0,
+        help="Only used when --kind adaptive_k.",
+    )
+    p.add_argument("--seed", type=int, default=0)
+    p.add_argument("--out_json", default="")
+    return p.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    set_seed(int(args.seed))
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    dtype = pick_dtype()
+
+    print(f"[eval-cellpolicy] kind={args.kind}  adapter={args.adapter_dir}", flush=True)
+    print(f"[eval-cellpolicy] eval_jsonl={args.eval_jsonl}  stage_i={args.stage_i}", flush=True)
+
+    tokenizer = AutoTokenizer.from_pretrained(
+        args.model_name, cache_dir=args.cache_dir, use_fast=True
+    )
+    if tokenizer.pad_token_id is None and tokenizer.eos_token_id is not None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    base = AutoModelForCausalLM.from_pretrained(
+        args.model_name, cache_dir=args.cache_dir, torch_dtype=dtype
+    )
+    model = PeftModel.from_pretrained(base, args.adapter_dir)
+    model.to(device)
+    model.eval()
+
+    rows: List[Dict[str, Any]] = load_jsonl_rows(args.eval_jsonl, limit_rows=int(args.eval_rows))
+    print(f"[eval-cellpolicy] loaded {len(rows)} eval rows", flush=True)
+
+    if args.kind == "strawman":
+        metrics = run_eval_strawman(
+            model, tokenizer, rows, device,
+            max_new_tokens=int(args.max_completion_length),
+            print_n=3,
+            stage_i=int(args.stage_i),
+        )
+    else:
+        metrics = run_eval_adaptive_k(
+            model, tokenizer, rows, device,
+            num_cot_tokens=int(args.num_cot_tokens),
+            max_new_tokens=int(args.max_completion_length),
+            print_n=3,
+            stage_i=int(args.stage_i),
+        )
+
+    print("[eval-cellpolicy] metrics:", json.dumps(metrics, indent=2), flush=True)
+    if args.out_json:
+        Path(args.out_json).parent.mkdir(parents=True, exist_ok=True)
+        with open(args.out_json, "w") as f:
+            json.dump(
+                {
+                    "kind": args.kind,
+                    "adapter_dir": args.adapter_dir,
+                    "eval_jsonl": args.eval_jsonl,
+                    "stage_i": int(args.stage_i),
+                    "num_cot_tokens": int(args.num_cot_tokens),
+                    "metrics": metrics,
+                },
+                f,
+                indent=2,
+            )
+        print(f"[eval-cellpolicy] wrote {args.out_json}", flush=True)
+
+
+if __name__ == "__main__":
+    main()

_runs/launch_adaptive_k_cellpolicy.sh

@@ -0,0 +1,42 @@
+#!/usr/bin/env bash
+# Launch two adaptive-k variants (single-stage cell-policy at stage_i=3,
+# no curriculum, but with growing recurrent-hidden thought tokens k).
+set -euo pipefail
+ROOT="${ROOT:-/home/ubuntu/curriculum_cot}"
+TS="$(date +%Y%m%d_%H%M%S)"
+SWEEP_ROOT="${ROOT}/_runs/adaptive_k_cellpolicy_${TS}"
+mkdir -p "${SWEEP_ROOT}"
+PY="${ROOT}/_runs/adaptive_k_cellpolicy_pipeline.py"
+
+launch() {
+  # Usage: launch <variant> <gpu> <KEY=VALUE>...  (ignored, args passed via positional CLI args)
+  local variant="$1" gpu="$2"
+  shift 2
+  local out="${SWEEP_ROOT}/${variant}"
+  mkdir -p "${out}"
+  echo "[launch] ${variant} on GPU ${gpu}  out=${out}"
+  nohup /opt/pytorch/bin/python -u "${PY}" \
+    --variant "${variant}" \
+    --gpu "${gpu}" \
+    --output_root "${out}" \
+    "$@" > "${out}/console.log" 2>&1 &
+  local pid=$!
+  disown "${pid}" || true
+  echo "${variant}=${pid}" >> "${SWEEP_ROOT}/PIDS.txt"
+}
+
+# adaptive_a: classic schedule (start at k=0, plateau-bumps with eps=0.01).
+launch adaptive_a_eps01 2 \
+  --start_k 0 --max_k 4 --steps_per_phase 600 --max_phases_per_k 2 \
+  --plateau_eps 0.01 --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
+
+# adaptive_b: faster k-growth (max_phases_per_k=1, force bump every phase).
+launch adaptive_b_fastgrow 3 \
+  --start_k 0 --max_k 4 --steps_per_phase 800 --max_phases_per_k 1 \
+  --plateau_eps 1.0 --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
+
+echo "[launch] sweep root: ${SWEEP_ROOT}"
+echo "[launch] PIDs:"
+cat "${SWEEP_ROOT}/PIDS.txt"

_runs/launch_adaptive_latent_baseline.sh

@@ -0,0 +1,76 @@
+#!/usr/bin/env bash
+# Adaptive-k thought-token baseline (experiment D in the 2x2 ablation).
+#
+# Same single-stage, whole-puzzle setup as launch_simple_baseline.sh
+# (experiment C, the "strawman"); same model, LoRA, JSONL, chat template.
+# The ONLY change is that the SFT loss uses the recurrent_hidden mechanism
+# with k thought tokens, and k grows automatically when the rolling-mean
+# loss plateaus.
+set -euo pipefail
+
+ROOT=/home/ubuntu/curriculum_cot
+SCRIPT=${ROOT}/_runs/adaptive_latent_baseline_sudoku_train.py
+PYTHON_BIN=/opt/pytorch/bin/python
+
+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
+
+SWEEP_ROOT=${ROOT}/_runs/adaptive_latent_$(date +%Y%m%d_%H%M%S)
+mkdir -p "${SWEEP_ROOT}"
+echo "${SWEEP_ROOT}" > "${ROOT}/_runs/current_adaptive_latent_sweep_dir"
+echo "SWEEP_ROOT=${SWEEP_ROOT}"
+
+export TOKENIZERS_PARALLELISM=false
+export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
+export HF_HOME="${ROOT}/.hf_cache"
+export TRANSFORMERS_CACHE="${ROOT}/.hf_cache"
+export WANDB_MODE=offline
+
+run_variant() {
+  local gpu="$1" tag="$2" lr="$3" max_k="$4" min_steps_per_k="$5"
+  local out=${SWEEP_ROOT}/${tag}
+  mkdir -p "${out}"
+  local log=${out}/train.log
+  : > "${log}"
+  (
+    export CUDA_VISIBLE_DEVICES="${gpu}"
+    "${PYTHON_BIN}" -u "${SCRIPT}" \
+      --train_jsonl "${TRAIN_JSONL}" \
+      --eval_jsonl "${EVAL_JSONL}" \
+      --output_dir "${out}" \
+      --learning_rate "${lr}" \
+      --max_steps 4000 \
+      --per_device_train_batch_size 4 \
+      --gradient_accumulation_steps 2 \
+      --logging_steps 25 \
+      --save_steps 500 \
+      --eval_every_steps 500 \
+      --eval_rows 50 \
+      --max_completion_length 96 \
+      --max_prompt_length 1024 \
+      --lora_r 32 --lora_alpha 64 --lora_dropout 0.05 \
+      --enable_gradient_checkpointing \
+      --start_k 0 \
+      --max_k "${max_k}" \
+      --min_steps_per_k "${min_steps_per_k}" \
+      --plateau_window 100 \
+      --plateau_eps 0.005 \
+      --converged_eps 0.001 \
+      --seed 0 \
+      >> "${log}" 2>&1
+  ) >/dev/null 2>&1 &
+  local pid=$!
+  echo "$pid $gpu $tag" >> "${SWEEP_ROOT}/PIDS.txt"
+  disown $pid 2>/dev/null || true
+  printf 'GPU %s -> %s pid=%s log=%s\n' "$gpu" "$tag" "$pid" "$log"
+}
+
+# 2 variants on idle GPUs 2,3:
+#  - adaptive_a:  same LR (5e-5) as strawman variant a, max_k=4, min_steps_per_k=400
+#  - adaptive_b:  smaller min_steps_per_k=250 to grow k more aggressively
+run_variant 2 adaptive_a_lr5e5_maxk4    5e-5  4  400
+run_variant 3 adaptive_b_lr5e5_fastgrow 5e-5  4  250
+
+echo
+echo "=== launched ==="
+cat "${SWEEP_ROOT}/PIDS.txt"

_runs/launch_baseline_1p5b_v4.sh

@@ -0,0 +1,82 @@
+#!/usr/bin/env bash
+# Launch 6 baseline 1.5B variants in parallel, one per GPU (0..5).
+# Each runs S2 GRPO -> S3 SFT -> S3 GRPO from a v3 lowsft S2 SFT checkpoint.
+set -euo pipefail
+
+ROOT="${ROOT:-/home/ubuntu/curriculum_cot}"
+SWEEP_ID="${SWEEP_ID:-$(date +%Y%m%d_%H%M%S)}"
+SWEEP_ROOT="${ROOT}/_runs/baseline_1p5b_v4_${SWEEP_ID}"
+PIPELINE="${ROOT}/_runs/baseline_1p5b_pipeline_v4.sh"
+
+mkdir -p "${SWEEP_ROOT}"
+SUMMARY="${SWEEP_ROOT}/SWEEP_README.md"
+
+CKPT_LR1E4="${ROOT}/checkpoints/sudoku-9x9-20empty-baseline-1p5b-sweep/baseline_lr1e4_lowsft_v3/s2_sft_v3/checkpoint-step-03000"
+CKPT_LR5E5="${ROOT}/checkpoints/sudoku-9x9-20empty-baseline-1p5b-sweep/baseline_lr5e5_lowsft_v3/s2_sft_v3/checkpoint-step-03000"
+
+if [[ ! -d "${CKPT_LR1E4}" || ! -d "${CKPT_LR5E5}" ]]; then
+  echo "ERROR: missing init checkpoints" >&2
+  exit 1
+fi
+
+cat >"${SUMMARY}" <<EOF
+# Baseline 1.5B v4 sweep — ${SWEEP_ID}
+
+Single GPU per variant. All 6 variants resume from the v3 lowsft S2 SFT
+checkpoints (the only ones with positive trend), then run S2 GRPO -> S3 SFT
+-> S3 GRPO with various GRPO LR / penalty / extra-S2-SFT settings.
+
+| GPU | variant | S2 init | GRPO LR | S3 SFT LR | penalty_singleton | extra S2 SFT (steps @ LR) |
+| ---: | --- | --- | ---: | ---: | ---: | --- |
+| 0 | pipe_a_lr1e4_grpo5e6        | lr1e4_lowsft step-3000 | 5e-6 | 2e-5 | 1.5 | 0 |
+| 1 | pipe_b_lr5e5_grpo5e6        | lr5e5_lowsft step-3000 | 5e-6 | 2e-5 | 1.5 | 0 |
+| 2 | pipe_c_lr1e4_grpo2e6        | lr1e4_lowsft step-3000 | 2e-6 | 2e-5 | 1.5 | 0 |
+| 3 | pipe_d_lr5e5_grpo2e6        | lr5e5_lowsft step-3000 | 2e-6 | 2e-5 | 1.5 | 0 |
+| 4 | pipe_e_lr5e5_grpo5e6_sngl25 | lr5e5_lowsft step-3000 | 5e-6 | 2e-5 | 2.5 | 0 |
+| 5 | pipe_f_lr1e4_extraS2sft     | lr1e4_lowsft step-3000 | 5e-6 | 2e-5 | 1.5 | 1500 @ 1e-5 |
+
+Pipeline budget per variant:
+- S2 GRPO  max 1200 steps (early stop on prec AND recall >= 0.98)
+- S3 SFT   max 2400 steps (same early stop)
+- S3 GRPO  max 1500 steps (same early stop)
+
+Logs: \`<variant>/PIPELINE.log\`, per-phase: \`<variant>/{s2_grpo,s3_sft,s3_grpo}/train.log\`
+EOF
+
+launch_variant() {
+  local gpu="$1" variant="$2" init="$3"
+  shift 3
+  local out="${SWEEP_ROOT}/${variant}"
+  mkdir -p "${out}"
+  local nohup_log="${out}/nohup.log"
+  printf 'GPU %s -> %s -> %s\n' "${gpu}" "${variant}" "${init}"
+  nohup env \
+    ROOT="${ROOT}" \
+    VARIANT="${variant}" \
+    GPU="${gpu}" \
+    S2_SFT_CKPT="${init}" \
+    OUTPUT_ROOT="${out}" \
+    USE_WANDB=0 \
+    WANDB_MODE=offline \
+    "$@" \
+    bash "${PIPELINE}" \
+    </dev/null >"${nohup_log}" 2>&1 &
+  local pid=$!
+  printf '  pid=%s   log=%s\n' "${pid}" "${nohup_log}"
+  echo "${pid} ${gpu} ${variant}" >> "${SWEEP_ROOT}/PIDS.txt"
+  disown "${pid}" 2>/dev/null || true
+}
+
+: > "${SWEEP_ROOT}/PIDS.txt"
+
+launch_variant 0 pipe_a_lr1e4_grpo5e6        "${CKPT_LR1E4}" GRPO_LR=5e-6 SFT_LR_S3=2e-5 PENALTY_SINGLETON=1.5
+launch_variant 1 pipe_b_lr5e5_grpo5e6        "${CKPT_LR5E5}" GRPO_LR=5e-6 SFT_LR_S3=2e-5 PENALTY_SINGLETON=1.5
+launch_variant 2 pipe_c_lr1e4_grpo2e6        "${CKPT_LR1E4}" GRPO_LR=2e-6 SFT_LR_S3=2e-5 PENALTY_SINGLETON=1.5
+launch_variant 3 pipe_d_lr5e5_grpo2e6        "${CKPT_LR5E5}" GRPO_LR=2e-6 SFT_LR_S3=2e-5 PENALTY_SINGLETON=1.5
+launch_variant 4 pipe_e_lr5e5_grpo5e6_sngl25 "${CKPT_LR5E5}" GRPO_LR=5e-6 SFT_LR_S3=2e-5 PENALTY_SINGLETON=2.5
+launch_variant 5 pipe_f_lr1e4_extraS2sft     "${CKPT_LR1E4}" GRPO_LR=5e-6 SFT_LR_S3=2e-5 PENALTY_SINGLETON=1.5 EXTRA_S2_SFT_STEPS=1500 EXTRA_S2_SFT_LR=1e-5
+
+echo
+echo "Sweep root: ${SWEEP_ROOT}"
+echo "Tail PIDS:"
+cat "${SWEEP_ROOT}/PIDS.txt"

_runs/launch_baseline_push_v5.sh

@@ -0,0 +1,84 @@
+#!/usr/bin/env bash
+# Wave-5: push baseline 1.5B past solve=0.35.
+#
+# Idea: best ckpts so far cap at per-cell-exact ~0.943 (solve 0.35 = 0.943^20).
+# To reach solve=0.5 we need exact ~= 0.965. That's +2.2pp of per-cell exact.
+#
+# 4 variants, single-GPU each, on GPUs 4..7.
+# All start from the leader (pipe_m post-S3-GRPO at solve=0.35) or its S3 SFT
+# ckpt, then push S3 GRPO further with different levers:
+#   - lower LR (escape / fine refine)
+#   - longer steps (3000 instead of 1500)
+#   - KL anchor (beta>0) to prevent regression
+#   - sharper rewards (mirror what worked for the latent's `s3_grpo_sharp_rwd`)
+set -euo pipefail
+
+ROOT=/home/ubuntu/curriculum_cot
+SWEEP_ROOT=/home/ubuntu/curriculum_cot/_runs/baseline_1p5b_v4_20260523_184952
+PIPELINE=$ROOT/_runs/baseline_1p5b_pipeline_v4.sh
+
+# best wave-2 anchors
+PIPE_M_S3GRPO_LATEST=$(ls -dt $SWEEP_ROOT/pipe_m_s3sft_from_b/s3_grpo/checkpoint-* 2>/dev/null | head -1)
+PIPE_M_S3SFT_LATEST=$SWEEP_ROOT/pipe_m_s3sft_from_b/s3_sft/checkpoint-step-02400
+PIPE_O_S3SFT_LATEST=$SWEEP_ROOT/pipe_o_s3sft_lr5e6/s3_sft/checkpoint-step-02400
+PIPE_J_S3GRPO_LATEST=$(ls -dt $SWEEP_ROOT/pipe_j_s3sft_lr5e5_lr1e5/s3_grpo/checkpoint-* 2>/dev/null | head -1)
+
+# Sanity
+for c in "$PIPE_M_S3GRPO_LATEST" "$PIPE_M_S3SFT_LATEST" "$PIPE_O_S3SFT_LATEST" "$PIPE_J_S3GRPO_LATEST"; do
+  [[ -d "$c" ]] || { echo "MISSING: $c"; exit 1; }
+done
+
+CKPT_LR5E5=$ROOT/checkpoints/sudoku-9x9-20empty-baseline-1p5b-sweep/baseline_lr5e5_lowsft_v3/s2_sft_v3/checkpoint-step-03000
+
+launch() {
+  local gpu="$1" variant="$2"; shift 2
+  local out=$SWEEP_ROOT/$variant; mkdir -p "$out"
+  nohup env ROOT="$ROOT" VARIANT="$variant" GPU="$gpu" S2_SFT_CKPT="$CKPT_LR5E5" \
+    OUTPUT_ROOT="$out" USE_WANDB=0 WANDB_MODE=offline "$@" \
+    bash "$PIPELINE" </dev/null >"$out/nohup.log" 2>&1 &
+  local pid=$!
+  echo "$pid $gpu $variant" >> "$SWEEP_ROOT/PIDS.txt"
+  disown $pid 2>/dev/null || true
+  printf 'GPU %s -> %s pid=%s\n' "$gpu" "$variant" "$pid"
+}
+
+# pipe_t (GPU 4): continue pipe_m's S3 GRPO with lower LR + KL anchor + longer steps.
+# Keep the policy near the SFT reference to avoid the regression we saw earlier.
+launch 4 pipe_t_grpo_low_kl \
+  START_PHASE=s3_grpo S3_GRPO_INIT="$PIPE_M_S3GRPO_LATEST" \
+  GRPO_LR=1e-6 GRPO_BS=32 GRPO_GA=1 GRPO_NG=8 \
+  GRPO_BETA=0.04 \
+  S3_GRPO_MAX_STEPS=3000 \
+  USE_GC=0
+
+# pipe_u (GPU 5): re-run S3 GRPO from pipe_m's S3-SFT ckpt with sharper rewards
+# (mirror latent `s3_grpo_sharp_rwd` recipe: bigger penalty for bad).
+launch 5 pipe_u_grpo_sharp_rwd \
+  START_PHASE=s3_grpo S3_GRPO_INIT="$PIPE_M_S3SFT_LATEST" \
+  GRPO_LR=5e-6 GRPO_BS=32 GRPO_GA=1 GRPO_NG=8 \
+  REWARD_GOOD=1.5 PENALTY_BAD=2.0 PENALTY_MAL=4.0 \
+  S3_GRPO_MAX_STEPS=3000 \
+  USE_GC=0
+
+# pipe_v (GPU 6): extend pipe_o's S3 SFT (the strongest pure-SFT path) with very
+# low LR for 4000 more steps. Then S3 GRPO at LR=1e-6.
+launch 6 pipe_v_sft_extend \
+  START_PHASE=s3_sft S3_SFT_INIT="$PIPE_O_S3SFT_LATEST" \
+  SFT_LR_S3=2e-6 SFT_BS=16 SFT_GA=1 \
+  S3_SFT_MAX_STEPS=4000 \
+  GRPO_LR=1e-6 GRPO_BS=32 GRPO_GA=1 GRPO_NG=8 \
+  S3_GRPO_MAX_STEPS=2000 \
+  USE_GC=0
+
+# pipe_w (GPU 7): continue pipe_j's S3 GRPO with very low LR + KL anchor.
+# Different lineage from pipe_m, so this gives an independent push.
+launch 7 pipe_w_j_low_kl \
+  START_PHASE=s3_grpo S3_GRPO_INIT="$PIPE_J_S3GRPO_LATEST" \
+  GRPO_LR=2e-6 GRPO_BS=32 GRPO_GA=1 GRPO_NG=8 \
+  GRPO_BETA=0.02 \
+  S3_GRPO_MAX_STEPS=3000 \
+  USE_GC=0
+
+echo
+echo "=== launched ==="
+cat "$SWEEP_ROOT/PIDS.txt" | tail -4

_runs/launch_baseline_push_v6.sh

@@ -0,0 +1,123 @@
+#!/usr/bin/env bash
+# Wave-6: push baseline 1.5B past solve=0.40 by porting the latent's winning
+# reward shaping + multi-value oversampling into the vanilla baseline pipeline.
+#
+# Diagnosis from v4/v5 logs:
+#   At plateau, eval shows `avg_set_size=1.000` for every step. The model
+#   is predicting only ONE value per cell even when the target is multi-valued.
+#   Per-cell exact pinned at 0.95 → solve = 0.95^20 ≈ 0.36.
+#   Same failure mode the latent's `s3_grpo_sharp_rwd` recipe fixed:
+#     exact_match_bonus + cardinality_mismatch_penalty + penalty_missing
+#   plus SFT-side multi_value_oversample_factor=5 (and target_size_min=2 for
+#   the most aggressive variant).
+#
+# 8 variants on GPUs 0..7. All seed from existing v4 best ckpts so we don't
+# burn cycles redoing S2.
+set -euo pipefail
+
+ROOT=/home/ubuntu/curriculum_cot
+SWEEP_ROOT=$ROOT/_runs/baseline_1p5b_v4_20260523_184952
+PIPELINE=$ROOT/_runs/baseline_1p5b_pipeline_v4.sh
+
+# --- v4 anchors ----
+PIPE_V_S3SFT_LATEST=$SWEEP_ROOT/pipe_v_sft_extend/s3_sft/checkpoint-step-04000
+PIPE_M_S3SFT_LATEST=$SWEEP_ROOT/pipe_m_s3sft_from_b/s3_sft/checkpoint-step-02400
+PIPE_V_S3GRPO_BEST=$SWEEP_ROOT/pipe_v_sft_extend/s3_grpo/checkpoint-1000   # step 1050 was 0.40 peak; 1000 is closest saved
+PIPE_M_S3GRPO_BEST=$SWEEP_ROOT/pipe_m_s3sft_from_b/s3_grpo/checkpoint-200  # peak per pipe_m logs
+PIPE_O_S3SFT_LATEST=$SWEEP_ROOT/pipe_o_s3sft_lr5e6/s3_sft/checkpoint-step-02400
+CKPT_LR5E5=$ROOT/checkpoints/sudoku-9x9-20empty-baseline-1p5b-sweep/baseline_lr5e5_lowsft_v3/s2_sft_v3/checkpoint-step-03000
+
+for c in "$PIPE_V_S3SFT_LATEST" "$PIPE_M_S3SFT_LATEST" "$PIPE_V_S3GRPO_BEST" "$PIPE_M_S3GRPO_BEST" "$PIPE_O_S3SFT_LATEST"; do
+  [[ -d "$c" ]] || { echo "MISSING: $c"; exit 1; }
+done
+
+launch() {
+  local gpu="$1" variant="$2"; shift 2
+  local out=$SWEEP_ROOT/$variant; mkdir -p "$out"
+  nohup env ROOT="$ROOT" VARIANT="$variant" GPU="$gpu" S2_SFT_CKPT="$CKPT_LR5E5" \
+    OUTPUT_ROOT="$out" USE_WANDB=0 WANDB_MODE=offline "$@" \
+    bash "$PIPELINE" </dev/null >"$out/nohup.log" 2>&1 &
+  local pid=$!
+  echo "$pid $gpu $variant" >> "$SWEEP_ROOT/PIDS.txt"
+  disown $pid 2>/dev/null || true
+  printf 'GPU %s -> %s pid=%s\n' "$gpu" "$variant" "$pid"
+}
+
+# === GRPO continuations (the high-leverage knob) ===
+
+# v6_a (GPU 0): continue best v4 GRPO with the FULL latent recipe.
+# card_pen=1.0 + missing=0.75 + exact_b=2.0; LR slightly lower than v4 to be safe.
+launch 0 v6_a_grpo_v_card \
+  START_PHASE=s3_grpo S3_GRPO_INIT="$PIPE_V_S3GRPO_BEST" \
+  GRPO_LR=2e-6 GRPO_BS=32 GRPO_GA=1 GRPO_NG=8 \
+  PENALTY_MISSING=0.75 EXACT_MATCH_BONUS=2.0 CARD_MISMATCH_PEN=1.0 \
+  S3_GRPO_MAX_STEPS=2000
+
+# v6_b (GPU 1): "sharp" version — mirror s3_grpo_sharp_rwd's stronger weights.
+launch 1 v6_b_grpo_v_sharp \
+  START_PHASE=s3_grpo S3_GRPO_INIT="$PIPE_V_S3GRPO_BEST" \
+  GRPO_LR=2e-6 GRPO_BS=32 GRPO_GA=1 GRPO_NG=8 \
+  PENALTY_MISSING=1.0 EXACT_MATCH_BONUS=4.0 CARD_MISMATCH_PEN=3.0 \
+  S3_GRPO_MAX_STEPS=2000
+
+# v6_c (GPU 2): full recipe but from pipe_v's S3 SFT (fresh GRPO, not continuation).
+launch 2 v6_c_grpo_vsft_card \
+  START_PHASE=s3_grpo S3_GRPO_INIT="$PIPE_V_S3SFT_LATEST" \
+  GRPO_LR=5e-6 GRPO_BS=32 GRPO_GA=1 GRPO_NG=8 \
+  PENALTY_MISSING=0.75 EXACT_MATCH_BONUS=2.0 CARD_MISMATCH_PEN=1.0 \
+  S3_GRPO_MAX_STEPS=2000
+
+# v6_d (GPU 3): same recipe but from pipe_m's S3 SFT (different lineage; champion).
+launch 3 v6_d_grpo_msft_card \
+  START_PHASE=s3_grpo S3_GRPO_INIT="$PIPE_M_S3SFT_LATEST" \
+  GRPO_LR=5e-6 GRPO_BS=32 GRPO_GA=1 GRPO_NG=8 \
+  PENALTY_MISSING=0.75 EXACT_MATCH_BONUS=2.0 CARD_MISMATCH_PEN=1.0 \
+  S3_GRPO_MAX_STEPS=2000
+
+# === SFT push w/ oversample (the data-side knob) ===
+
+# v6_e (GPU 4): continue pipe_v S3 SFT with oversample=5. Mirrors r1_sft_c_oversample5.
+launch 4 v6_e_sft_v_oversample5 \
+  START_PHASE=s3_sft S3_SFT_INIT="$PIPE_V_S3SFT_LATEST" \
+  SFT_LR_S3=2e-6 SFT_BS=16 SFT_GA=1 \
+  SFT_OVERSAMPLE=5 \
+  S3_SFT_MAX_STEPS=2500 \
+  GRPO_LR=2e-6 GRPO_BS=32 GRPO_GA=1 GRPO_NG=8 \
+  PENALTY_MISSING=0.75 EXACT_MATCH_BONUS=2.0 CARD_MISMATCH_PEN=1.0 \
+  S3_GRPO_MAX_STEPS=1500
+
+# v6_f (GPU 5): same but oversample=8 (more aggressive).
+launch 5 v6_f_sft_v_oversample8 \
+  START_PHASE=s3_sft S3_SFT_INIT="$PIPE_V_S3SFT_LATEST" \
+  SFT_LR_S3=2e-6 SFT_BS=16 SFT_GA=1 \
+  SFT_OVERSAMPLE=8 \
+  S3_SFT_MAX_STEPS=2500 \
+  GRPO_LR=2e-6 GRPO_BS=32 GRPO_GA=1 GRPO_NG=8 \
+  PENALTY_MISSING=0.75 EXACT_MATCH_BONUS=2.0 CARD_MISMATCH_PEN=1.0 \
+  S3_GRPO_MAX_STEPS=1500
+
+# v6_g (GPU 6): oversample=5 + train_target_size_min=2 (only multi-value cells).
+# This is the most surgical variant — focus all training mass on the failing cells.
+launch 6 v6_g_sft_v_mv_only \
+  START_PHASE=s3_sft S3_SFT_INIT="$PIPE_V_S3SFT_LATEST" \
+  SFT_LR_S3=1e-6 SFT_BS=16 SFT_GA=1 \
+  SFT_OVERSAMPLE=5 SFT_TGT_MIN=2 \
+  S3_SFT_MAX_STEPS=2000 \
+  GRPO_LR=2e-6 GRPO_BS=32 GRPO_GA=1 GRPO_NG=8 \
+  PENALTY_MISSING=0.75 EXACT_MATCH_BONUS=2.0 CARD_MISMATCH_PEN=1.0 \
+  S3_GRPO_MAX_STEPS=1500
+
+# v6_h (GPU 7): same as v6_a but with even more steps + KL anchor for stability.
+# The latent best (s3_grpo_baseline) ran with beta=0.0; we know KL>0 hurts long term.
+# But here we want to see whether the new shape rewards survive more steps without
+# regression. Use a small beta (0.01) for gentle anchoring.
+launch 7 v6_h_grpo_v_card_long \
+  START_PHASE=s3_grpo S3_GRPO_INIT="$PIPE_V_S3GRPO_BEST" \
+  GRPO_LR=2e-6 GRPO_BS=32 GRPO_GA=1 GRPO_NG=8 \
+  GRPO_BETA=0.01 \
+  PENALTY_MISSING=0.75 EXACT_MATCH_BONUS=2.0 CARD_MISMATCH_PEN=1.0 \
+  S3_GRPO_MAX_STEPS=3000
+
+echo
+echo "=== launched ==="
+tail -8 "$SWEEP_ROOT/PIDS.txt"

_runs/launch_latent_reproduction_overnight.sh

@@ -0,0 +1,82 @@
+#!/usr/bin/env bash
+# Overnight reproduction of the latent recurrent-hidden 3-stage champion.
+# Mirrors the recipe that produced solve=0.60 (100p) / 0.675 (40p) on 2026-05-22.
+#
+# Single distributed job across all 8 H100s. End-to-end runtime: ~6-7 hrs.
+#
+# Stages: S1 SFT (cot=1) -> S1 GRPO (cot=1)
+#       -> S2 SFT (cot=2) -> S2 GRPO (cot=2)
+#       -> S3 SFT (cot=3) -> S3 GRPO (cot=3)
+#
+# Hyperparameters (defaults, faithful to original):
+#   model  Qwen/Qwen2.5-0.5B-Instruct
+#   num_cot_tokens 1->2->3 across stages
+#   latent_mode recurrent_hidden
+#   bs=8/device, grad_accum=2, gradient checkpointing ON
+#   stage1_sft_lr=2e-4, stage2/3_sft_lr=5e-5, grpo_lr=1e-6 (hardcoded)
+#   value_target=0.98 (precision AND recall)
+#   train_puzzles=10000  eval_puzzles=100
+#   num_generations=4  max_completion_length=24
+
+set -euo pipefail
+
+ROOT=/home/ubuntu/curriculum_cot
+SCRIPT="${ROOT}/hard_9x9_stage1_consistency_queue/launch_20empty_latent_recurrent_stages123_value98.sh"
+
+RUN_TAG="latent_reproduction_overnight_$(date +%Y%m%d_%H%M%S)"
+OUTPUT_ROOT="${ROOT}/_runs/${RUN_TAG}"
+LOG="${OUTPUT_ROOT}/PIPELINE.log"
+mkdir -p "${OUTPUT_ROOT}"
+
+# Free the HF caches and ensure our pre-downloaded Qwen 0.5B is found
+export HF_HOME="${ROOT}/.hf_cache"
+export TRANSFORMERS_CACHE="${ROOT}/.hf_cache"
+export HF_HUB_OFFLINE=0
+export TOKENIZERS_PARALLELISM=false
+export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
+# wandb is not authenticated on this machine — keep offline so jobs don't hang
+export WANDB_MODE=offline
+# avoid the upstream script trying to pull from a wandb entity we don't own
+export WANDB_ENTITY="local"
+
+# Use our preinstalled pytorch venv
+export PYTHON_BIN=/opt/pytorch/bin/python
+
+# 4-GPU distributed run with doubled grad accum to preserve the original
+# effective batch size (8*2*8 = 128 -> 8*4*4 = 128). Takes ~2x wall-clock
+# but is faithful to the original convergence dynamics.
+export GPU_IDS=0,1,2,3
+export NUM_PROCESSES=4
+export SFT_GRAD_ACCUM=4
+export GRPO_GRAD_ACCUM=4
+
+# Match original
+export MODEL_NAME="Qwen/Qwen2.5-0.5B-Instruct"
+export VALUE_TARGET=0.98
+export SFT_VALUE_TARGET=0.95
+export GRPO_VALUE_TARGET=0.98
+export TRAIN_PUZZLES=10000
+export EVAL_PUZZLES=100
+export MIN_STEPS_BEFORE_STOP=50
+
+# Cap per-phase wallclock to keep us safely under one overnight session.
+# The original took ~6-7 hours; we cap each phase at 75 min to let all 6 phases
+# finish within ~7.5 hrs even if one phase slow-runs.
+export PHASE_WALL_CLOCK_SECONDS=4500
+
+# Hard step caps (in addition to early stop on prec+recall)
+export SFT_MAX_STEPS=4000
+export GRPO_MAX_STEPS=2000
+
+export RUN_TAG
+export OUTPUT_ROOT
+export CHECKPOINT_ROOT="${OUTPUT_ROOT}"
+
+printf '[launch_latent_reproduction] %s\n' "$(date -Is)" | tee -a "${LOG}"
+printf '  RUN_TAG=%s\n' "${RUN_TAG}" | tee -a "${LOG}"
+printf '  OUTPUT_ROOT=%s\n' "${OUTPUT_ROOT}" | tee -a "${LOG}"
+printf '  GPUs=%s nproc=%s model=%s\n' "${GPU_IDS}" "${NUM_PROCESSES}" "${MODEL_NAME}" | tee -a "${LOG}"
+printf '  VALUE_TARGET=%s SFT_VALUE_TARGET=%s GRPO_VALUE_TARGET=%s\n' "${VALUE_TARGET}" "${SFT_VALUE_TARGET}" "${GRPO_VALUE_TARGET}" | tee -a "${LOG}"
+printf '  PHASE_WALL_CLOCK=%ss SFT_MAX_STEPS=%s GRPO_MAX_STEPS=%s\n' "${PHASE_WALL_CLOCK_SECONDS}" "${SFT_MAX_STEPS}" "${GRPO_MAX_STEPS}" | tee -a "${LOG}"
+
+bash "${SCRIPT}" 2>&1 | tee -a "${LOG}"

_runs/launch_simple_baseline.sh

@@ -0,0 +1,97 @@
+#!/usr/bin/env bash
+# Strawman baseline for the rebuttal: vanilla LoRA, no curriculum, no thought
+# tokens, single-shot whole-puzzle prediction. SFT followed by GRPO.
+#
+# Same model (Qwen2.5-1.5B-Instruct), same LoRA (r=32, α=64, dropout=0.05),
+# same JSONL data file, same Qwen chat template wrapping as the cell-policy
+# experiments. The ONLY differences from the cell-policy baseline are:
+#   - no per-cell expansion (one example per puzzle)
+#   - no stage_i / curriculum
+#   - no multi_value_oversample, no exact_match_bonus / cardinality penalties
+#   - reward = number of correct values out of 20 + whole-solve bonus
+set -euo pipefail
+
+ROOT=/home/ubuntu/curriculum_cot
+SCRIPT=${ROOT}/_runs/simple_baseline_sudoku_train.py
+PYTHON_BIN=/opt/pytorch/bin/python
+
+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
+
+SWEEP_ROOT=${ROOT}/_runs/strawman_baseline_$(date +%Y%m%d_%H%M%S)
+mkdir -p "${SWEEP_ROOT}"
+echo "${SWEEP_ROOT}" > "${ROOT}/_runs/current_strawman_sweep_dir"
+echo "SWEEP_ROOT=${SWEEP_ROOT}"
+
+export TOKENIZERS_PARALLELISM=false
+export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
+export HF_HOME="${ROOT}/.hf_cache"
+export TRANSFORMERS_CACHE="${ROOT}/.hf_cache"
+export WANDB_MODE=offline
+
+run_pipeline() {
+  local gpu="$1" tag="$2" sft_lr="$3" grpo_lr="$4" sft_max="$5" grpo_max="$6"
+  local out=${SWEEP_ROOT}/${tag}
+  mkdir -p "${out}"
+  local log=${out}/pipeline.log
+  : > "${log}"
+  (
+    export CUDA_VISIBLE_DEVICES="${gpu}"
+    echo "[$(date +%H:%M:%S)] === ${tag} on GPU ${gpu}: SFT lr=${sft_lr} max_steps=${sft_max} ===" >> "${log}"
+    "${PYTHON_BIN}" -u "${SCRIPT}" \
+      --phase sft \
+      --train_jsonl "${TRAIN_JSONL}" \
+      --eval_jsonl "${EVAL_JSONL}" \
+      --output_dir "${out}/sft" \
+      --learning_rate "${sft_lr}" \
+      --max_steps "${sft_max}" \
+      --per_device_train_batch_size 8 \
+      --gradient_accumulation_steps 2 \
+      --num_epochs 8 \
+      --logging_steps 25 \
+      --save_steps 200 \
+      --eval_rows 100 \
+      --max_completion_length 96 \
+      --max_prompt_length 1024 \
+      --lora_r 32 --lora_alpha 64 --lora_dropout 0.05 \
+      --seed 0 \
+      >> "${log}" 2>&1
+
+    echo "[$(date +%H:%M:%S)] === ${tag} on GPU ${gpu}: GRPO lr=${grpo_lr} max_steps=${grpo_max} ===" >> "${log}"
+    "${PYTHON_BIN}" -u "${SCRIPT}" \
+      --phase grpo \
+      --init_adapter_dir "${out}/sft/final" \
+      --train_jsonl "${TRAIN_JSONL}" \
+      --eval_jsonl "${EVAL_JSONL}" \
+      --output_dir "${out}/grpo" \
+      --learning_rate "${grpo_lr}" \
+      --max_steps "${grpo_max}" \
+      --per_device_train_batch_size 4 \
+      --gradient_accumulation_steps 2 \
+      --num_generations 8 \
+      --beta 0.0 \
+      --temperature 1.0 \
+      --num_epochs 50 \
+      --logging_steps 25 \
+      --save_steps 200 \
+      --eval_rows 100 \
+      --max_completion_length 96 \
+      --max_prompt_length 1024 \
+      --lora_r 32 --lora_alpha 64 --lora_dropout 0.05 \
+      --seed 0 \
+      >> "${log}" 2>&1
+    echo "[$(date +%H:%M:%S)] === ${tag} DONE ===" >> "${log}"
+  ) >/dev/null 2>&1 &
+  local pid=$!
+  echo "$pid $gpu $tag" >> "${SWEEP_ROOT}/PIDS.txt"
+  disown $pid 2>/dev/null || true
+  printf 'GPU %s -> %s pid=%s log=%s\n' "$gpu" "$tag" "$pid" "$log"
+}
+
+# 2 variants on GPUs 0,1: explore SFT LR (5e-5 and 1e-4) — same GRPO LR (5e-6).
+run_pipeline 0 strawman_a_sft5e5_grpo5e6  5e-5  5e-6  2000  1500
+run_pipeline 1 strawman_b_sft1e4_grpo5e6  1e-4  5e-6  2000  1500
+
+echo
+echo "=== launched ==="
+cat "${SWEEP_ROOT}/PIDS.txt"

_runs/launch_strawman_cellpolicy.sh

@@ -0,0 +1,38 @@
+#!/usr/bin/env bash
+# Launch two strawman variants (single-stage cell-policy at stage_i=3, no
+# curriculum, no thought tokens) on GPUs 0 and 1.
+set -euo pipefail
+ROOT="${ROOT:-/home/ubuntu/curriculum_cot}"
+TS="$(date +%Y%m%d_%H%M%S)"
+SWEEP_ROOT="${ROOT}/_runs/strawman_cellpolicy_${TS}"
+mkdir -p "${SWEEP_ROOT}"
+PIPE="${ROOT}/_runs/strawman_cellpolicy_pipeline.sh"
+chmod +x "${PIPE}"
+
+launch() {
+  # Usage: launch <variant> <gpu> <KEY=VALUE>...
+  local variant="$1" gpu="$2"
+  shift 2
+  local out="${SWEEP_ROOT}/${variant}"
+  mkdir -p "${out}"
+  echo "[launch] ${variant} on GPU ${gpu}  out=${out}"
+  nohup env VARIANT="${variant}" GPU="${gpu}" OUTPUT_ROOT="${out}" "$@" \
+       bash "${PIPE}" > "${out}/console.log" 2>&1 &
+  local pid=$!
+  disown "${pid}" || true
+  echo "${variant}=${pid}" >> "${SWEEP_ROOT}/PIDS.txt"
+}
+
+launch strawman_a_lr2e5 0 \
+  SFT_LR=2e-5 GRPO_LR=5e-6 SFT_MAX_STEPS=3000 GRPO_MAX_STEPS=1500 \
+  PENALTY_MISSING=1.0 EXACT_MATCH_BONUS=1.0 CARD_MISMATCH_PEN=1.5 \
+  SFT_OVERSAMPLE=3
+
+launch strawman_b_lr5e5 1 \
+  SFT_LR=5e-5 GRPO_LR=5e-6 SFT_MAX_STEPS=4000 GRPO_MAX_STEPS=1500 \
+  PENALTY_MISSING=1.0 EXACT_MATCH_BONUS=1.0 CARD_MISMATCH_PEN=1.5 \
+  SFT_OVERSAMPLE=3
+
+echo "[launch] sweep root: ${SWEEP_ROOT}"
+echo "[launch] PIDs:"
+cat "${SWEEP_ROOT}/PIDS.txt"

_runs/simple_baseline_sudoku_train.py

@@ -0,0 +1,559 @@
+#!/usr/bin/env python3
+"""Strawman baseline for the rebuttal.
+
+Vanilla Qwen2.5-1.5B-Instruct + LoRA on top of the *existing* JSONL data
+(`data/sudoku_t3_20empty_value_qwen_text_stage1_{train,eval}.jsonl`).
+
+Compared to the cell-policy / latent recipes, this strawman intentionally
+removes everything that helped:
+
+- NO curriculum (single stage; we don't even read `stage_i`).
+- NO chain-of-thought / latent thought tokens.
+- NO per-cell expansion (one example == one whole puzzle).
+- NO multi-value oversampling, no special reward shaping (just matches/N).
+
+It uses the *same* model, *same* LoRA config, *same* tokenizer + chat
+template wrapping that every cell-policy experiment used, so any solve
+gap vs the cell-policy / latent runs is purely due to task framing,
+not data, prompt, model, or PEFT differences.
+
+Usage:
+    python simple_baseline_sudoku_train.py --phase sft  --output_dir <out>/sft  --learning_rate 5e-5
+    python simple_baseline_sudoku_train.py --phase grpo --init_adapter_dir <out>/sft/final \
+        --output_dir <out>/grpo --learning_rate 5e-6
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import math
+import os
+import re
+import sys
+import time
+from pathlib import Path
+from typing import Any, Callable, Dict, List, Optional
+
+import torch
+from datasets import Dataset
+from peft import LoraConfig, PeftModel, get_peft_model
+from transformers import AutoModelForCausalLM, AutoTokenizer, set_seed
+
+# Reuse existing helpers (these are the canonical ones used by every cell-policy run).
+ROOT = Path(__file__).resolve().parent.parent
+if str(ROOT) not in sys.path:
+    sys.path.insert(0, str(ROOT))
+
+from multi_output_cell_policy.sft_multi_output_train import (  # type: ignore
+    load_jsonl_rows,
+    pick_dtype,
+)
+from multi_output_cell_policy.rewards import score_prediction_text  # type: ignore
+from multi_output_cell_policy.shared_multi_output_policy import (  # type: ignore
+    make_solved_grid_from_row,
+    stage_i_consistent_values,
+)
+from aligned_cell_policy.shared_cell_policy import build_cell_examples_from_row  # type: ignore
+
+
+# ---- Strawman task definition -----------------------------------------------
+# This is the ONLY new piece relative to the cell-policy experiments. The
+# system prompt asks the model to emit the missing values for ALL empty cells
+# in one shot, in the row-major order that the existing JSONL `completion`
+# field already uses. The user message is the raw `prompt` field from the
+# JSONL (puzzle as (row,col,value) tuples), which is byte-identical to what
+# `prompt_builder.py` consumes in cell-policy runs.
+
+SYSTEM_PROMPT_STRAWMAN = (
+    "You are a Sudoku solver.\n"
+    "You will be given a 9x9 Sudoku grid encoded as (row,col,value) tuples in "
+    "row-major order, where value 0 marks an empty cell.\n"
+    "Predict the missing values for ALL empty cells in row-major order.\n"
+    "Return ONLY a JSON list of integers like [v1,v2,...,vK], where K is the "
+    "number of empty cells (typically 20). Each value must be an integer in "
+    "[1,9].\n"
+    "Do not include any explanation, markdown, or text outside the JSON list."
+)
+
+
+def build_chat_prompt(tokenizer: Any, raw_prompt: str) -> str:
+    """Same chat template wrapping every other experiment uses (Qwen, system+user)."""
+    messages = [
+        {"role": "system", "content": SYSTEM_PROMPT_STRAWMAN.strip()},
+        {"role": "user", "content": raw_prompt},
+    ]
+    chat_template = getattr(tokenizer, "chat_template", None)
+    if chat_template:
+        return tokenizer.apply_chat_template(
+            messages, tokenize=False, add_generation_prompt=True
+        )
+    return SYSTEM_PROMPT_STRAWMAN.strip() + "\n\n" + raw_prompt + "\n"
+
+
+# ---- Reward -----------------------------------------------------------------
+
+LIST_RE = re.compile(r"\[[^\[\]]*\]")
+
+
+def parse_int_list(text: str) -> Optional[List[int]]:
+    """Parse the model's emission as a JSON int list with values in [1,9].
+
+    Tolerant: tries the whole completion first, then falls back to the first
+    well-formed JSON list match. Returns None on failure.
+    """
+    s = str(text).strip()
+    if not s:
+        return None
+    candidates: List[str] = []
+    candidates.append(s)
+    m = LIST_RE.search(s)
+    if m is not None:
+        candidates.append(m.group(0))
+    for cand in candidates:
+        try:
+            obj = json.loads(cand)
+        except Exception:
+            continue
+        if not isinstance(obj, list):
+            continue
+        out: List[int] = []
+        ok = True
+        for v in obj:
+            if isinstance(v, bool) or not isinstance(v, int):
+                ok = False
+                break
+            if v < 1 or v > 9:
+                ok = False
+                break
+            out.append(int(v))
+        if ok:
+            return out
+    return None
+
+
+def whole_puzzle_reward(
+    *,
+    pred_list: Optional[List[int]],
+    target_list: List[int],
+    parse_penalty: float = 4.0,
+    length_mismatch_penalty: float = 0.5,
+    full_solve_bonus: float = 5.0,
+) -> float:
+    """Simple reward: matches per cell + bonus for full solve, penalty if parse fails."""
+    if pred_list is None:
+        return -float(parse_penalty)
+    n = len(target_list)
+    matches = 0
+    for i in range(min(len(pred_list), n)):
+        if int(pred_list[i]) == int(target_list[i]):
+            matches += 1
+    reward = float(matches)
+    if len(pred_list) != n:
+        reward -= float(length_mismatch_penalty) * abs(len(pred_list) - n)
+    if len(pred_list) == n and matches == n:
+        reward += float(full_solve_bonus)
+    return reward
+
+
+# ---- Dataset construction ---------------------------------------------------
+
+
+def build_dataset(rows: List[Dict[str, Any]], tokenizer: Any) -> Dataset:
+    prompts, completions, targets = [], [], []
+    for row in rows:
+        raw_prompt = str(row["prompt"]).strip()
+        completion_str = str(row["completion"]).strip()
+        target = parse_int_list(completion_str)
+        if target is None:
+            continue
+        prompts.append(build_chat_prompt(tokenizer, raw_prompt))
+        completions.append(completion_str)
+        targets.append(json.dumps(target, separators=(",", ":")))
+    return Dataset.from_dict(
+        {"prompt": prompts, "completion": completions, "target": targets}
+    )
+
+
+# ---- Eval (deterministic, greedy, single-shot) ------------------------------
+
+
+@torch.no_grad()
+@torch.no_grad()
+def run_eval(
+    model: torch.nn.Module,
+    tokenizer: Any,
+    eval_rows: List[Dict[str, Any]],
+    device: torch.device,
+    max_new_tokens: int = 96,
+    print_n: int = 3,
+    stage_i: int = 3,
+) -> Dict[str, float]:
+    """Apples-to-apples eval with the cell-policy framework.
+
+    The strawman model emits the WHOLE puzzle (a JSON list of integers) in
+    one forward pass. We then split that list into per-cell SINGLETON
+    predictions and score each cell with the same ``score_prediction_text``
+    function the cell-policy / latent baselines use, against the i-consistent
+    target set at ``stage_i`` (default 3 — matching the S3 eval used for the
+    rebuttal v6 baseline and the latent champion).
+
+    Reported metrics mirror ``multi_output_cell_policy/sft_multi_output_train.py::run_eval``
+    so numbers are directly comparable across all four 2x2 ablation cells.
+    """
+    model.eval()
+    total_cells = 0
+    parse_ok = 0.0
+    canonical_ok = 0.0
+    exact_set_match = 0.0
+    includes_gt = 0.0
+    precision_sum = 0.0
+    recall_sum = 0.0
+    cardinality_match_sum = 0.0
+    n_solve = 0
+    n_total_puzzles = 0
+    n_parse_fail_puzzles = 0
+    printed = 0
+    for row in eval_rows:
+        target_completion = parse_int_list(str(row["completion"]))
+        if target_completion is None:
+            continue
+        n_total_puzzles += 1
+        prompt = build_chat_prompt(tokenizer, str(row["prompt"]).strip())
+        enc = tokenizer(prompt, return_tensors="pt", add_special_tokens=False)
+        enc = {k: v.to(device) for k, v in enc.items()}
+        out = model.generate(
+            **enc,
+            max_new_tokens=int(max_new_tokens),
+            do_sample=False,
+            eos_token_id=tokenizer.eos_token_id,
+            pad_token_id=tokenizer.pad_token_id,
+        )
+        gen = tokenizer.decode(
+            out[0][int(enc["input_ids"].shape[1]) :], skip_special_tokens=True
+        ).strip()
+        pred_list = parse_int_list(gen)
+
+        try:
+            cells = build_cell_examples_from_row(row)
+            solved = make_solved_grid_from_row(row)
+        except Exception as e:
+            if printed < print_n:
+                print(f"[strawman eval debug] row skipped (no metadata): {e}", flush=True)
+                printed += 1
+            continue
+
+        row_all_exact = True
+        row_has_eval_cell = False
+        for idx, ex in enumerate(cells):
+            target_values = stage_i_consistent_values(
+                ex.grid, target_cell=ex.target_cell, stage_i=int(stage_i)
+            )
+            row_has_eval_cell = True
+            if pred_list is not None and idx < len(pred_list):
+                pred_text = json.dumps({"values": [int(pred_list[idx])]})
+            else:
+                pred_text = ""
+            info = score_prediction_text(
+                text=pred_text,
+                grid=ex.grid,
+                solved=solved,
+                target_cell=ex.target_cell,
+                stage_i=int(stage_i),
+                reward_good_value=1.0,
+                penalty_bad_value=1.75,
+                penalty_malformed=4.0,
+                penalty_empty=0.5,
+                penalty_singleton=1.5,
+            )
+            total_cells += 1
+            parse_ok += float(info["parse_ok"])
+            canonical_ok += float(info["strict_canonical"])
+            exact_set_match += float(info["exact_set_match"])
+            includes_gt += float(info["includes_ground_truth"])
+            precision_sum += float(info["value_precision"])
+            recall_sum += float(info["value_recall"])
+            if int(info["num_predicted_values"]) == int(len(target_values)):
+                cardinality_match_sum += 1.0
+            if float(info["exact_set_match"]) < 0.5:
+                row_all_exact = False
+        if row_has_eval_cell and row_all_exact:
+            n_solve += 1
+        if pred_list is None:
+            n_parse_fail_puzzles += 1
+        if printed < print_n:
+            head_pred = pred_list if pred_list is not None else "PARSE_FAIL"
+            print(
+                f"[strawman eval debug] target={target_completion}  pred={head_pred}  "
+                f"solve={int(row_all_exact and row_has_eval_cell)}  gen={gen!r}",
+                flush=True,
+            )
+            printed += 1
+    return {
+        "n_total_cells": float(total_cells),
+        "n_total_puzzles": float(n_total_puzzles),
+        "parse_rate": float(parse_ok / max(1, total_cells)),
+        "strict_canonical_rate": float(canonical_ok / max(1, total_cells)),
+        "exact_set_match_rate": float(exact_set_match / max(1, total_cells)),
+        "includes_ground_truth_rate": float(includes_gt / max(1, total_cells)),
+        "value_precision": float(precision_sum / max(1, total_cells)),
+        "value_recall": float(recall_sum / max(1, total_cells)),
+        "cardinality_match_rate": float(cardinality_match_sum / max(1, total_cells)),
+        "puzzle_parse_fail_rate": float(n_parse_fail_puzzles / max(1, n_total_puzzles)),
+        "solve_rate": float(n_solve) / max(1, n_total_puzzles),
+    }
+
+
+# ---- Main -------------------------------------------------------------------
+
+
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser()
+    p.add_argument("--phase", choices=["sft", "grpo"], required=True)
+    p.add_argument("--model_name", type=str, default="Qwen/Qwen2.5-1.5B-Instruct")
+    p.add_argument("--train_jsonl", type=str, required=True)
+    p.add_argument("--eval_jsonl", type=str, required=True)
+    p.add_argument("--output_dir", type=str, required=True)
+    p.add_argument("--cache_dir", type=str, default=str(ROOT / ".hf_cache"))
+    p.add_argument("--init_adapter_dir", type=str, default="")
+    p.add_argument("--seed", type=int, default=0)
+
+    # Data
+    p.add_argument("--limit_train_rows", type=int, default=10000)
+    p.add_argument("--eval_rows", type=int, default=100)
+
+    # Train hyperparameters
+    p.add_argument("--per_device_train_batch_size", type=int, default=8)
+    p.add_argument("--gradient_accumulation_steps", type=int, default=2)
+    p.add_argument("--learning_rate", type=float, default=5e-5)
+    p.add_argument("--weight_decay", type=float, default=0.0)
+    p.add_argument("--num_epochs", type=float, default=8.0)
+    p.add_argument("--max_steps", type=int, default=2000)
+    p.add_argument("--logging_steps", type=int, default=25)
+    p.add_argument("--save_steps", type=int, default=200)
+    p.add_argument("--eval_steps", type=int, default=150)
+    p.add_argument("--max_grad_norm", type=float, default=1.0)
+    p.add_argument("--max_completion_length", type=int, default=96)
+    p.add_argument("--max_prompt_length", type=int, default=1024)
+
+    # LoRA
+    p.add_argument("--lora_r", type=int, default=32)
+    p.add_argument("--lora_alpha", type=int, default=64)
+    p.add_argument("--lora_dropout", type=float, default=0.05)
+    p.add_argument("--enable_gradient_checkpointing", action="store_true")
+
+    # GRPO-only
+    p.add_argument("--num_generations", type=int, default=8)
+    p.add_argument("--beta", type=float, default=0.0)
+    p.add_argument("--temperature", type=float, default=1.0)
+    p.add_argument("--full_solve_bonus", type=float, default=5.0)
+    p.add_argument("--length_mismatch_penalty", type=float, default=0.5)
+    p.add_argument("--parse_penalty", type=float, default=4.0)
+
+    # W&B
+    p.add_argument("--use_wandb", action="store_true")
+    p.add_argument("--wandb_project", type=str, default="sudoku-strawman-baseline")
+    p.add_argument("--wandb_run_name", type=str, default="")
+    p.add_argument("--wandb_mode", type=str, default="offline")
+
+    return p.parse_args()
+
+
+def setup_model_and_tokenizer(args: argparse.Namespace, device: torch.device):
+    tokenizer = AutoTokenizer.from_pretrained(
+        args.model_name, cache_dir=args.cache_dir, use_fast=True
+    )
+    if tokenizer.pad_token_id is None:
+        tokenizer.pad_token = tokenizer.eos_token or "<|endoftext|>"
+    if tokenizer.padding_side != "left":
+        tokenizer.padding_side = "left"
+
+    model = AutoModelForCausalLM.from_pretrained(
+        args.model_name,
+        cache_dir=args.cache_dir,
+        torch_dtype=pick_dtype(),
+        low_cpu_mem_usage=True,
+    )
+    if str(args.init_adapter_dir).strip():
+        model = PeftModel.from_pretrained(model, args.init_adapter_dir, is_trainable=True)
+    else:
+        lora = LoraConfig(
+            r=args.lora_r,
+            lora_alpha=args.lora_alpha,
+            lora_dropout=args.lora_dropout,
+            bias="none",
+            task_type="CAUSAL_LM",
+            target_modules=[
+                "q_proj", "k_proj", "v_proj", "o_proj",
+                "gate_proj", "up_proj", "down_proj",
+            ],
+        )
+        model = get_peft_model(model, lora)
+
+    if args.enable_gradient_checkpointing:
+        if hasattr(model, "gradient_checkpointing_enable"):
+            model.gradient_checkpointing_enable(
+                gradient_checkpointing_kwargs={"use_reentrant": False}
+            )
+        if hasattr(model, "enable_input_require_grads"):
+            model.enable_input_require_grads()
+    if hasattr(model, "config"):
+        model.config.use_cache = False
+    model.to(device)
+    return model, tokenizer
+
+
+def run_sft(args: argparse.Namespace) -> None:
+    from trl import SFTConfig, SFTTrainer  # type: ignore
+
+    set_seed(int(args.seed))
+    os.makedirs(args.output_dir, exist_ok=True)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    train_rows = load_jsonl_rows(args.train_jsonl, limit_rows=int(args.limit_train_rows))
+    eval_rows = load_jsonl_rows(args.eval_jsonl, limit_rows=int(args.eval_rows))
+
+    model, tokenizer = setup_model_and_tokenizer(args, device)
+
+    # Build dataset of {prompt, completion} where prompt is chat-templated.
+    train_ds = build_dataset(train_rows, tokenizer)
+
+    cfg = SFTConfig(
+        output_dir=args.output_dir,
+        per_device_train_batch_size=int(args.per_device_train_batch_size),
+        gradient_accumulation_steps=int(args.gradient_accumulation_steps),
+        learning_rate=float(args.learning_rate),
+        weight_decay=float(args.weight_decay),
+        num_train_epochs=float(args.num_epochs),
+        max_steps=int(args.max_steps),
+        logging_steps=int(args.logging_steps),
+        save_steps=int(args.save_steps),
+        save_strategy="steps",
+        save_total_limit=4,
+        eval_strategy="no",
+        bf16=(pick_dtype() == torch.bfloat16),
+        fp16=(pick_dtype() == torch.float16),
+        max_grad_norm=float(args.max_grad_norm),
+        gradient_checkpointing=bool(args.enable_gradient_checkpointing),
+        report_to=("wandb" if args.use_wandb else "none"),
+        run_name=(args.wandb_run_name or None),
+        max_length=int(args.max_prompt_length + args.max_completion_length + 8),
+        completion_only_loss=True,
+        seed=int(args.seed),
+    )
+
+    trainer = SFTTrainer(
+        model=model,
+        args=cfg,
+        train_dataset=train_ds,
+        processing_class=tokenizer,
+    )
+
+    # Periodic eval hook (TRL doesn't natively give us a custom eval loop hook,
+    # so we run eval before training and after the final step here).
+    print("[strawman sft] BEFORE-train eval:", run_eval(model, tokenizer, eval_rows, device), flush=True)
+
+    t0 = time.time()
+    trainer.train()
+    print(f"[strawman sft] training time = {time.time() - t0:.1f}s", flush=True)
+
+    final_dir = os.path.join(args.output_dir, "final")
+    trainer.save_model(final_dir)
+    print(f"[strawman sft] saved final adapter to {final_dir}", flush=True)
+
+    print("[strawman sft] AFTER-train eval:", run_eval(model, tokenizer, eval_rows, device), flush=True)
+
+
+def run_grpo(args: argparse.Namespace) -> None:
+    from trl import GRPOConfig, GRPOTrainer  # type: ignore
+
+    set_seed(int(args.seed))
+    os.makedirs(args.output_dir, exist_ok=True)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    train_rows = load_jsonl_rows(args.train_jsonl, limit_rows=int(args.limit_train_rows))
+    eval_rows = load_jsonl_rows(args.eval_jsonl, limit_rows=int(args.eval_rows))
+
+    model, tokenizer = setup_model_and_tokenizer(args, device)
+    train_ds = build_dataset(train_rows, tokenizer)
+
+    parse_penalty = float(args.parse_penalty)
+    length_mismatch_penalty = float(args.length_mismatch_penalty)
+    full_solve_bonus = float(args.full_solve_bonus)
+
+    def reward_fn(completions, target, **kwargs):
+        rewards: List[float] = []
+        for c, tgt in zip(completions, target):
+            tgt_list = json.loads(tgt) if isinstance(tgt, str) else list(tgt)
+            pred = parse_int_list(str(c))
+            rewards.append(
+                whole_puzzle_reward(
+                    pred_list=pred,
+                    target_list=tgt_list,
+                    parse_penalty=parse_penalty,
+                    length_mismatch_penalty=length_mismatch_penalty,
+                    full_solve_bonus=full_solve_bonus,
+                )
+            )
+        return rewards
+
+    cfg = GRPOConfig(
+        output_dir=args.output_dir,
+        per_device_train_batch_size=int(args.per_device_train_batch_size),
+        gradient_accumulation_steps=int(args.gradient_accumulation_steps),
+        learning_rate=float(args.learning_rate),
+        weight_decay=float(args.weight_decay),
+        num_train_epochs=float(args.num_epochs),
+        max_steps=int(args.max_steps),
+        logging_steps=int(args.logging_steps),
+        save_steps=int(args.save_steps),
+        save_strategy="steps",
+        save_total_limit=6,
+        bf16=(pick_dtype() == torch.bfloat16),
+        fp16=(pick_dtype() == torch.float16),
+        max_grad_norm=float(args.max_grad_norm),
+        gradient_checkpointing=bool(args.enable_gradient_checkpointing),
+        report_to=("wandb" if args.use_wandb else "none"),
+        run_name=(args.wandb_run_name or None),
+        max_prompt_length=int(args.max_prompt_length),
+        max_completion_length=int(args.max_completion_length),
+        num_generations=int(args.num_generations),
+        beta=float(args.beta),
+        temperature=float(args.temperature),
+        seed=int(args.seed),
+    )
+
+    trainer = GRPOTrainer(
+        model=model,
+        reward_funcs=[reward_fn],
+        args=cfg,
+        train_dataset=train_ds,
+        processing_class=tokenizer,
+    )
+
+    print("[strawman grpo] BEFORE-train eval:", run_eval(model, tokenizer, eval_rows, device), flush=True)
+
+    t0 = time.time()
+    trainer.train()
+    print(f"[strawman grpo] training time = {time.time() - t0:.1f}s", flush=True)
+
+    final_dir = os.path.join(args.output_dir, "final")
+    trainer.save_model(final_dir)
+    print(f"[strawman grpo] saved final adapter to {final_dir}", flush=True)
+
+    print("[strawman grpo] AFTER-train eval:", run_eval(model, tokenizer, eval_rows, device), flush=True)
+
+
+def main() -> None:
+    args = parse_args()
+    if args.use_wandb:
+        os.environ.setdefault("WANDB_MODE", str(args.wandb_mode))
+        os.environ["WANDB_PROJECT"] = args.wandb_project
+    if args.phase == "sft":
+        run_sft(args)
+    else:
+        run_grpo(args)
+
+
+if __name__ == "__main__":
+    main()

_runs/status.sh

@@ -0,0 +1,42 @@
+#!/usr/bin/env bash
+# One-shot snapshot of the active sweep.
+SWEEP="${1:-$(ls -dt /home/ubuntu/curriculum_cot/_runs/baseline_1p5b_v4_* 2>/dev/null | head -1)}"
+[[ -z "${SWEEP}" || ! -d "${SWEEP}" ]] && { echo "no sweep"; exit 1; }
+echo "=== sweep: ${SWEEP} ==="
+echo "=== nvidia-smi ==="
+nvidia-smi --query-gpu=index,utilization.gpu,memory.used,memory.total,power.draw --format=csv,noheader
+echo
+echo "=== pids ==="
+while read -r pid gpu name; do
+  if kill -0 "$pid" 2>/dev/null; then alive=ALIVE; else alive=DEAD; fi
+  printf '  pid=%-6s gpu=%s %-30s %s\n' "$pid" "$gpu" "$name" "$alive"
+done < "${SWEEP}/PIDS.txt"
+echo
+echo "=== per-variant phase + best/last eval ==="
+for v in "${SWEEP}"/pipe_*; do
+  vn="$(basename "$v")"
+  current_phase="(starting)"
+  for ph in s2_sft_extra s2_grpo s3_sft s3_grpo; do
+    [[ -d "$v/$ph" ]] && current_phase="$ph"
+  done
+  printf '\n--- %s (phase=%s) ---\n' "$vn" "${current_phase}"
+  # Pipeline log tail
+  if [[ -f "$v/PIPELINE.log" ]]; then
+    tail -3 "$v/PIPELINE.log" | sed 's/^/    PL: /'
+  fi
+  # Phase-specific evals
+  for ph in s2_sft_extra s2_grpo s3_sft s3_grpo; do
+    log="$v/$ph/train.log"
+    [[ -f "$log" ]] || continue
+    # SFT eval lines
+    last_sft="$(grep -E "\[baseline sft eval\] " "$log" 2>/dev/null | tail -3)"
+    last_grpo="$(grep -E "\[baseline grpo (custom )?eval" "$log" 2>/dev/null | tail -3)"
+    last_train="$(grep -E "\[baseline (sft|grpo) (train|final)" "$log" 2>/dev/null | tail -1)"
+    if [[ -n "$last_sft$last_grpo$last_train" ]]; then
+      printf '  [%s]\n' "$ph"
+      [[ -n "$last_train" ]] && echo "$last_train" | sed 's/^/      tr: /'
+      [[ -n "$last_sft" ]]   && echo "$last_sft"   | sed 's/^/      ev: /'
+      [[ -n "$last_grpo" ]]  && echo "$last_grpo"  | sed 's/^/      ev: /'
+    fi
+  done
+done

_runs/strawman_cellpolicy_pipeline.sh

@@ -0,0 +1,186 @@
+#!/usr/bin/env bash
+# Strawman = single-stage cell-policy at stage_i=3 from BASE (no curriculum,
+# no thought tokens). Same per-cell prompt, same trainer scripts, same scoring
+# function as the v6 baseline and the latent champion. The ONLY differences
+# vs the v6 baseline are:
+#   - No prior SFT/GRPO at stage_i=1 or stage_i=2 (start fresh from base Qwen).
+#   - Single SFT phase + single GRPO phase, both at stage_i=3.
+#   - No latent recurrent-hidden tokens (vanilla LoRA on base model).
+# Required env vars: VARIANT, GPU, OUTPUT_ROOT.
+set -euo pipefail
+
+ROOT="${ROOT:-/home/ubuntu/curriculum_cot}"
+PYTHON_BIN="${PYTHON_BIN:-/opt/pytorch/bin/python}"
+SFT_SCRIPT="${ROOT}/multi_output_cell_policy/sft_multi_output_train.py"
+GRPO_SCRIPT="${ROOT}/multi_output_cell_policy/grpo_multi_output_train.py"
+
+: "${VARIANT:?VARIANT required}"
+: "${GPU:?GPU required}"
+
+OUTPUT_ROOT="${OUTPUT_ROOT:-${ROOT}/_runs/strawman_cellpolicy_$(date +%Y%m%d_%H%M%S)/${VARIANT}}"
+MODEL_NAME="${MODEL_NAME:-Qwen/Qwen2.5-1.5B-Instruct}"
+
+# Use the same S3 hyperparameters as the v6 baseline so the only knob is
+# "did we do the curriculum or not".
+SFT_LR="${SFT_LR:-2e-5}"
+SFT_BS="${SFT_BS:-16}"
+SFT_GA="${SFT_GA:-2}"
+SFT_MAX_STEPS="${SFT_MAX_STEPS:-3000}"
+
+GRPO_LR="${GRPO_LR:-5e-6}"
+GRPO_BETA="${GRPO_BETA:-0.0}"
+GRPO_NG="${GRPO_NG:-8}"
+GRPO_BS="${GRPO_BS:-16}"
+GRPO_GA="${GRPO_GA:-2}"
+GRPO_PROMPT="${GRPO_PROMPT:-768}"
+GRPO_COMPL="${GRPO_COMPL:-24}"
+GRPO_MAX_STEPS="${GRPO_MAX_STEPS:-1500}"
+
+# v6-style reward shaping (same as the v6 sweep that hit solve=0.44).
+REWARD_GOOD="${REWARD_GOOD:-1.25}"
+PENALTY_BAD="${PENALTY_BAD:-1.0}"
+PENALTY_MAL="${PENALTY_MAL:-4.0}"
+PENALTY_EMPTY="${PENALTY_EMPTY:-0.5}"
+PENALTY_SINGLETON="${PENALTY_SINGLETON:-1.5}"
+PENALTY_MISSING="${PENALTY_MISSING:-1.0}"
+EXACT_MATCH_BONUS="${EXACT_MATCH_BONUS:-1.0}"
+CARD_MISMATCH_PEN="${CARD_MISMATCH_PEN:-1.5}"
+SFT_OVERSAMPLE="${SFT_OVERSAMPLE:-3}"
+SFT_TGT_MIN="${SFT_TGT_MIN:-0}"
+SFT_TGT_MAX="${SFT_TGT_MAX:-0}"
+
+VALUE_TARGET="${VALUE_TARGET:-0.98}"
+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}"
+
+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"
+
+mkdir -p "${OUTPUT_ROOT}"
+PIPELINE_LOG="${OUTPUT_ROOT}/PIPELINE.log"
+ts() { date +'%H:%M:%S'; }
+log() { printf '[%s] %s\n' "$(ts)" "$*" | tee -a "${PIPELINE_LOG}" >&2; }
+
+best_ckpt() {
+  local d="$1"
+  if [[ -f "${d}/adapter_model.safetensors" ]]; then
+    printf '%s\n' "${d}"; return 0
+  fi
+  shopt -s nullglob
+  local cks=("${d}"/checkpoint-step-* "${d}"/checkpoint-*)
+  shopt -u nullglob
+  (( ${#cks[@]} == 0 )) && return 1
+  printf '%s\n' "${cks[@]}" | sort -V | tail -n 1
+}
+
+if [[ ! -f "${TRAIN_JSONL}" || ! -f "${EVAL_JSONL}" ]]; then
+  log "ERROR: missing dataset jsonls"; exit 1
+fi
+
+export CUDA_VISIBLE_DEVICES="${GPU}"
+export TOKENIZERS_PARALLELISM=false
+export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
+export HF_HOME="${ROOT}/.hf_cache"
+export TRANSFORMERS_CACHE="${ROOT}/.hf_cache"
+
+GC_FLAG=()
+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 "    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}"
+
+# ----- 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) ==="
+"${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 "" \
+  --seed 0 \
+  --gpu_id 0 \
+  --stage_i 3 \
+  --total_empties_hint 20 \
+  --per_device_train_batch_size "${SFT_BS}" \
+  --gradient_accumulation_steps "${SFT_GA}" \
+  --num_epochs 256 \
+  --learning_rate "${SFT_LR}" \
+  --max_grad_norm 1.0 \
+  --logging_steps 25 \
+  --eval_steps 200 \
+  --save_steps 200 \
+  --eval_rows "${EVAL_ROWS}" \
+  --max_completion_length 24 \
+  --limit_train_rows "${TRAIN_ROWS}" \
+  --lora_r 32 --lora_alpha 64 --lora_dropout 0.05 \
+  --eval_value_precision_stop "${VALUE_TARGET}" \
+  --eval_value_recall_stop "${VALUE_TARGET}" \
+  --eval_exact_set_match_stop 0 \
+  --eval_solve_rate_stop 0 \
+  --min_steps_before_stop 200 \
+  --max_wall_clock_seconds "${PHASE_WALL_SECS}" \
+  --max_steps "${SFT_MAX_STEPS}" \
+  --multi_value_oversample_factor "${SFT_OVERSAMPLE}" \
+  --train_target_size_min "${SFT_TGT_MIN}" \
+  --train_target_size_max "${SFT_TGT_MAX}" \
+  "${GC_FLAG[@]}" 2>&1 | tee "${SFT_DIR}/train.log"
+
+SFT_CKPT="$(best_ckpt "${SFT_DIR}")" || { log "ERROR: no SFT ckpt"; exit 1; }
+log ">>> SFT ckpt: ${SFT_CKPT}"
+
+# ----- Phase 2: GRPO at stage_i=3 from SFT output -----
+GRPO_DIR="${OUTPUT_ROOT}/grpo"
+mkdir -p "${GRPO_DIR}"
+log "=== PHASE GRPO (stage_i=3, init=${SFT_CKPT}) ==="
+"${PYTHON_BIN}" -u "${GRPO_SCRIPT}" \
+  --model_name "${MODEL_NAME}" \
+  --train_jsonl "${TRAIN_JSONL}" \
+  --eval_jsonl "${EVAL_JSONL}" \
+  --output_dir "${GRPO_DIR}" \
+  --cache_dir "${ROOT}/.hf_cache" \
+  --init_adapter_dir "${SFT_CKPT}" \
+  --seed 0 \
+  --gpu_id 0 \
+  --stage_i 3 \
+  --total_empties_hint 20 \
+  --per_device_train_batch_size "${GRPO_BS}" \
+  --gradient_accumulation_steps "${GRPO_GA}" \
+  --num_train_epochs 100 \
+  --learning_rate "${GRPO_LR}" \
+  --logging_steps 10 \
+  --save_steps 200 \
+  --eval_steps 150 \
+  --eval_rows "${EVAL_ROWS}" \
+  --num_generations "${GRPO_NG}" \
+  --max_prompt_length "${GRPO_PROMPT}" \
+  --max_completion_length "${GRPO_COMPL}" \
+  --beta "${GRPO_BETA}" \
+  --limit_train_rows "${TRAIN_ROWS}" \
+  --lora_r 32 --lora_alpha 64 --lora_dropout 0.05 \
+  --reward_good_value "${REWARD_GOOD}" \
+  --penalty_bad_value "${PENALTY_BAD}" \
+  --penalty_malformed "${PENALTY_MAL}" \
+  --penalty_empty "${PENALTY_EMPTY}" \
+  --penalty_singleton "${PENALTY_SINGLETON}" \
+  --penalty_missing "${PENALTY_MISSING}" \
+  --exact_match_bonus "${EXACT_MATCH_BONUS}" \
+  --cardinality_mismatch_penalty "${CARD_MISMATCH_PEN}" \
+  --eval_value_precision_stop "${VALUE_TARGET}" \
+  --eval_value_recall_stop "${VALUE_TARGET}" \
+  --eval_solve_rate_stop 0 \
+  --min_steps_before_stop 100 \
+  --max_wall_clock_seconds "${PHASE_WALL_SECS}" \
+  --max_steps "${GRPO_MAX_STEPS}" \
+  "${GC_FLAG[@]}" 2>&1 | tee "${GRPO_DIR}/train.log"
+
+GRPO_CKPT="$(best_ckpt "${GRPO_DIR}")" || { log "WARN: no GRPO ckpt found"; exit 0; }
+log ">>> GRPO ckpt: ${GRPO_CKPT}"
+log "===== STRAWMAN ${VARIANT} done ====="