Blame for code/docs/K5_LEVERS.md

# K5: inference-time levers at matched generated tokens

Status: harness built and CPU-tested (2026-10-19); no GPU run yet. Plan: `docs/SOTA_PLAN.md` v3.1 section 3 (K5, "run FIRST"),
`docs/PREREG_SOTA.md` section 4 (K5 MDE 1.3 pp estimator test) and section 5 (K5 is its own Holm family).

Question: at the SAME total generated tokens per problem, does Lean-feedback correction (L1) or the agentic loop (L6: proof
bank + APOLLO-style subgoal repair + compiler feedback) solve more dev problems than restart sampling?

## 1. What is built

| File | Role |
|---|---|
| `scripts/p2/k5_levers.py` | Orchestrator: stage plan (exact command lines), resumable runner, per-problem token ledger, budget equaliser, proof bank + decontamination + retrieval, route/assemble rounds of the agentic arm, matched-budget report with paired tests |
| `align/rl/envs/subgoal_repair.py` | Pure-Python subgoal extraction / reassembly; Lean only through `align.verify.lean` (read-only: `Worker.header_env` / `Worker.command`, the same calls as `verify_generations.collect_feedback`) |
| `extras/vllm_rollout/lean_prompts.py`, `lean_gen.py` | New prompt format `goedel_bank` (Goedel prompt preceded by retrieved verified proofs; with no examples it is byte-identical to `goedel`). The `goedel` path is untouched (its byte-pin test passes). |
| `tests/test_k5_levers.py`, `tests/test_subgoal_repair.py`, `tests/fixtures/k5_baseline_proofs.jsonl` | 30 + 22 tests; fixture = 14 failed and 6 accepted real proofs from the P2 baseline file, verified with `--feedback` |

## 2. Arms and the matching rule

All arms: Goedel-Prover-V2-8B (rev dfd02e62), T = 1.0, top_p = 0.95, seed 0, cap 40,960, same dev problems, same frozen verifier.

* **restart**: n independent round-0 attempts per problem (also the roots of `correct`, so the arms are paired and the roots are generated once).
* **correct** (class B): each failed root gets up to 2 Goedel-format repairs (`lean_gen --round 1/2 --corrections-from`, `--trajectory-cap-tokens 40000`, `--repair-solved` so every trajectory has a complete cost record, default `--no-feedback-policy skip`).
* **agentic** (class C, exploratory), R = 4 rounds: round 0 = bank-conditioned attempt (`goedel_bank`, k = 3 retrieved proofs); rounds 1..3 = `subgoal`, `feedback`, `subgoal` (configurable `--routes`; a design choice, not tuned). A `subgoal` round falls back to `feedback` for parents with nothing extractable.

**Unit and cost.** The unit of every arm is a trajectory. Its cost is the completion tokens of every node up to the end of the round of its first success (all nodes of a round are paid, branches included), or all its tokens if it never succeeded. Subgoal samples (accepted or not) are charged to the trajectory that spawned them, through the assembled node's `trajectory.completion_tokens`. Prompt/prefill tokens are not in the budget (same convention as the class-B ledger); they are reported separately in box-hours.

**Equaliser.** Arms are generated generously as pools of units, then replayed at a per-problem budget B_p: units are drawn in random order (256 permutations, seed fixed per problem) until one succeeds or the next unit does not fit the remaining budget (a real run cannot know a unit's cost in advance; it would be clipped and fail, so the replay stops rather than skipping to a cheaper unit). The per-problem estimate is the fraction of orderings that solve it. B_p defaults to `restart_k:8` = 8 x the mean restart attempt tokens of that problem (the spend of a pass@8 run); `fixed:N` is available. Curves at 0.25 / 0.5 / 1.0 x B_p are reported. A problem whose pool cannot fill the budget is `underfilled`; `k5_levers.py status` prints how many more units to generate. Contrasts (correct - restart, agentic - restart, agentic - correct) use `align.eval.paired.sign_flip_test`, `paired_bootstrap_ci` and `holm` (the K5 family of PREREG section 5). The K0 training-seed null does not apply (no training).

Tested arithmetic (tests/test_k5_levers.py): closed-form replay (4 equal units, 1 solved, budget for 2 -> 0.5; for 3 -> 0.75), never overspends, deterministic, stop-rule, underfill planning, `restart_k` budget, a 2x-costlier chain loses to restarts at equal tokens, chain/branch/assembled-node cost, refusal of unverified or partial lineages.

## 3. Subgoal repair (APOLLO-style)

1. **Locate** (`locate_targets`): every Lean error (positions index `compile_view(code)`, translated by `view_line_map`) maps to the innermost failing `have name : T := rhs` (indent/bracket based extents, masked comments), else a main-line tactic, else a "tail" (proof ended with goals). Errors on a `have` header are reported unrepairable. A target containing a finer target is dropped (the outer failure shows up next round). More than 6 gaps: refused.
2. **Sorrify** (`build_skeleton`): replace each target by `sorry`.
3. **Goal states**: one REPL command on the skeleton; its `sorries` give the exact hypotheses and goal at each gap. The tail cut is added only if the skeleton still reports the unsolved-goals error at the theorem's `by` (after a parse error Lean reports a cascade that vanishes once the broken step is cut; found on a real case).
4. **Lemma** (`lemma_from_goal`): `theorem k5_sg_<hash> (hyps) : goal := by sorry`; inaccessible hypotheses (`n✝`) get fresh names with every later mention rewritten, instances become `[inst]`; let-hyps unsupported. Each lemma is checked to elaborate (`check_lemma_elaborates`); those that do not are not sent to the prover. They are ordinary problems for `lean_gen` (Goedel round-0 prompt + bank examples).
5. **Reassemble** (`reassemble`): the prover's tactic body is spliced inline at the gap's indentation, preceded by `rename_i` for inaccessible hypotheses so the body sees the lemma's names. Unproven gaps get their original text back (so a half-repaired proof is a better parent next round). Inline rather than lemma-first, because the strict verifier rejects extra declarations unless `allow_helper_defs`, which we do not change.
6. The assembled proof is only a candidate: it is verified by the unchanged strict verifier against the ORIGINAL statement.

Real-data evidence (all pass): every one of 14 real failures yields targets; skeleton + restore is the identity on all 14; 36 real `have` cuts of 6 accepted proofs re-inlined reproduce the proof up to whitespace; under real Lean (`-m slow`): goal states returned for 8 failures, at least 5 lemmas elaborate, restoring gaps gives the same REPL errors as the original, and the full cut -> lemma proven by the original body (strict-accepted) -> splice -> strict-accepted loop succeeds on 5 real proofs and on a crafted case that needs `rename_i`. In an exploratory run on the 14 failures, the first gap's lemma failed to elaborate for 2 (a `ZMod` numeral ambiguity, which stays a limitation, and a shadowed `n✝`, fixed by the rename rewrite).

## 4. Proof bank and decontamination

Source: accepted round-0 proofs of the TRAIN pool (`train_pool.jsonl`, `ver_r0.frozen.jsonl`), the shortest per problem (<= 2,500 chars). Decontamination (`decontaminate_bank`) uses `align.data.decontam` (normalised-text, signature and alpha hashes, MinHash-LSH, token n-gram containment, informal-text n-grams) against EVERY evaluation and query set passed (`--eval-problems` plus the dev problems themselves), drops ids equal to an eval id, and propagates ancestry. Over-flagging is the policy; dropped entries and reasons are written to `bank_dropped.jsonl`. Retrieval is TF-IDF cosine over statement tokens (identifiers kept), excluding the query's own problem. Tests: copies, renames, hypothesis reorderings and eval-id collisions are dropped; a near-copy of a dev problem is dropped; on the real pool, a fresh detector finds zero hits in the kept set and the kept set shares no id with any eval set. NOT yet done: the elaborated-type hash (needs the Lean env; `scripts/p2/supply/elab_hash.py`) is supported by the detector but not fed here, so reordered-hypothesis copies rely on MinHash/n-grams. Subgoal-lemma prompts retrieve from the same bank.

## 5. Design sources

* APOLLO (Ospanov, Farnia, Yousefzadeh), arXiv 2505.05758: Lean compiler + LLM; repair syntax, isolate failing sub-lemmas, automated solvers, small-budget LLM calls on remaining goals, recombine and recheck. We take: failing-step isolation, per-subgoal proving, reassembly, re-verification. We do not take: the automated-solver pass (hammer tactics) or syntax-refiner agent; adding them would confound the "model" lever.
* Hilbert (Varambally et al.), arXiv 2509.22819 (ICLR 2026): recursive decomposition of unsolved subgoals, verifier feedback, semantic theorem retrieval. We take: retrieval-conditioned proving and the idea that unsolved subgoals recurse (our per-round re-extraction on the assembled proof). We do not take the informal-reasoner model (needs a second, stronger model; breaks "same 8B").
* DeepSeek-Prover-V2, arXiv 2504.21801: subgoal decomposition with `have ... := sorry` skeletons, prove subgoals separately, recombine. Our skeleton-of-sorries is the same idea, built from errors instead of a decomposing model. (Cited from memory; not re-fetched in this session.)
* Goedel-Prover-V2, arXiv 2508.03613: the correction prompt and the 2-repair protocol, ported byte-faithfully in `lean_prompts.py`. (From memory.)
* Kimina-Prover / Kimina Lean Server, arXiv 2504.11354: REPL-pool verification, which `align.verify.lean` mirrors. (From memory.)
* OProver-8B (the plan's class-C bar: feedback ablation 91.8 vs 87.7, SR): motivation for the bank + feedback combination; no code consulted.

## 6. Limits (read before interpreting)

* The agentic arm confounds three levers (bank, subgoal repair, feedback). Attribution needs ablations (bank off via empty `bank_examples`; `--routes feedback,feedback,feedback`), not run here.
* Bank-conditioned prompts are a different conditioning of the policy than the Goedel prompt; `goedel_bank` is not covered by the Goedel byte-fidelity gate.
* Assembled nodes carry placeholder token ids (`completion_ids = [0]`); their real cost is in `trajectory.completion_tokens`. Anything that trains on these records must exclude `kind == "assembled"`.
* The replay estimator treats units as exchangeable draws from the generated pool; correlated units (same root) are not an issue because each root is one unit, but pool size caps what budget is testable (`underfilled`).
* The class-B arm uses Goedel's skip policy for truncated or error-less parents, so truncated roots are wasted; agentic uses the same.
* Prefill/verifier CPU cost is not in the token budget; box-hours below include generation only.
* Not run on GPU: `route`/`assemble` have CPU tests (assemble end-to-end with fake probe, route only via the unit pieces); `route` against a live pool is covered by the slow `plan_subgoals` tests, not by a CLI test.

## 7. GPU runs (dev set = miniF2F-valid + ProofNet-valid), estimated box-hours

Prerequisites: `proofnet_valid.jsonl` / `proofnet_test.jsonl` built (Phase 0a item 3; not present under `.tmp/p2/problems/` at the time of writing), chain8 step `pool` finished (`r1/pool/ver_r0.frozen.jsonl`), and the verifier frozen (`.tmp/restart/VERIFIER_FROZEN`).

```
cd /apps/development/rlhf/align && . /apps/development/rlhf/env/apps_only.sh
P=.tmp/p2; K5=$P/k5/dev
PYTHONPATH=$PWD .venv/bin/python scripts/p2/k5_levers.py run --run-dir $K5 \
  --problems $P/problems/minif2f_valid.jsonl $P/problems/proofnet_valid.jsonl \
  --bank-problems $P/problems/train_pool.jsonl --bank-gens $P/r1/pool/gen_r0.jsonl --bank-vers $P/r1/pool/ver_r0.frozen.jsonl \
  --eval-problems $P/problems/minif2f.jsonl $P/problems/proofnet_test.jsonl $P/problems/mobench.jsonl \
  --n-restart 16 --n-agentic 6 --budget restart_k:8
```
`run` executes the 18 stages below in order, each inside `hostguard.py run --wait` (GPU stages 24 cores / 100 GB, CPU stages 16 cores / 64 GB), each with a done-marker (`$K5/done/<stage>`), so a re-run resumes; every generation/verification file is keyed and resumable as before. `plan` with the same arguments prints all 18 expanded lines without starting anything; `--only <stage>...` runs a subset. The first stages, expanded (paths real, $K5 = `.tmp/p2/k5/dev`):

```
# restart_r0
/apps/development/rlhf/align/.venv/bin/python /apps/development/rlhf/align/scripts/hostguard.py run --wait --cores 24 --mem-gb 100 --label k5-restart_r0 -- /apps/development/rlhf/align/.venv/bin/python /apps/development/rlhf/align/scripts/p2/pipeline_gen_verify.py --run-dir /apps/development/rlhf/align/.tmp/p2/k5/dev/pipe/restart_r0 --problems /apps/development/rlhf/align/.tmp/p2/k5/dev/dev.jsonl --verify-out /apps/development/rlhf/align/.tmp/p2/k5/dev/restart/ver_r0.jsonl --workers 8 --label k5-restart_r0 --gen-args '--model /apps/development/rlhf/hf_cache/hub/models--Goedel-LM--Goedel-Prover-V2-8B/snapshots/dfd02e6271a58375dfbf3ece0175277cf6b6a89a --problems /apps/development/rlhf/align/.tmp/p2/k5/dev/dev.jsonl --out /apps/development/rlhf/align/.tmp/p2/k5/dev/restart/gen_r0.jsonl --n 32 --cap-tokens 40960 --round 0 --tp 4 --seed 0 --max-num-seqs 128 --gpu-mem-util 0.82 --label k5-restart_r0' --feedback

# correct_r1
/apps/development/rlhf/align/.venv/bin/python /apps/development/rlhf/align/scripts/hostguard.py run --wait --cores 24 --mem-gb 100 --label k5-correct_r1 -- /apps/development/rlhf/align/.venv/bin/python /apps/development/rlhf/align/scripts/p2/pipeline_gen_verify.py --run-dir /apps/development/rlhf/align/.tmp/p2/k5/dev/pipe/correct_r1 --problems /apps/development/rlhf/align/.tmp/p2/k5/dev/dev.jsonl --verify-out /apps/development/rlhf/align/.tmp/p2/k5/dev/correct/ver_r1.jsonl --workers 8 --label k5-correct_r1 --gen-args '--model /apps/development/rlhf/hf_cache/hub/models--Goedel-LM--Goedel-Prover-V2-8B/snapshots/dfd02e6271a58375dfbf3ece0175277cf6b6a89a --problems /apps/development/rlhf/align/.tmp/p2/k5/dev/dev.jsonl --out /apps/development/rlhf/align/.tmp/p2/k5/dev/correct/gen_r1.jsonl --n 1 --cap-tokens 40960 --round 1 --tp 4 --seed 0 --max-num-seqs 128 --gpu-mem-util 0.82 --label k5-correct_r1 --trajectory-cap-tokens 40000 --repair-solved --corrections-from /apps/development/rlhf/align/.tmp/p2/k5/dev/restart/ver_r0.jsonl /apps/development/rlhf/align/.tmp/p2/k5/dev/restart/gen_r0.jsonl' --feedback

# agentic_r0
/apps/development/rlhf/align/.venv/bin/python /apps/development/rlhf/align/scripts/hostguard.py run --wait --cores 24 --mem-gb 100 --label k5-agentic_r0 -- /apps/development/rlhf/align/.venv/bin/python /apps/development/rlhf/align/scripts/p2/pipeline_gen_verify.py --run-dir /apps/development/rlhf/align/.tmp/p2/k5/dev/pipe/agentic_r0 --problems /apps/development/rlhf/align/.tmp/p2/k5/dev/dev_bank.jsonl --verify-out /apps/development/rlhf/align/.tmp/p2/k5/dev/agentic/ver_r0.jsonl --workers 8 --label k5-agentic_r0 --gen-args '--model /apps/development/rlhf/hf_cache/hub/models--Goedel-LM--Goedel-Prover-V2-8B/snapshots/dfd02e6271a58375dfbf3ece0175277cf6b6a89a --problems /apps/development/rlhf/align/.tmp/p2/k5/dev/dev_bank.jsonl --out /apps/development/rlhf/align/.tmp/p2/k5/dev/agentic/gen_r0.jsonl --n 8 --cap-tokens 40960 --round 0 --tp 4 --seed 0 --max-num-seqs 128 --gpu-mem-util 0.82 --label k5-agentic_r0 --prompt-format goedel_bank' --feedback

# agentic_route_r1
/apps/development/rlhf/align/.venv/bin/python /apps/development/rlhf/align/scripts/hostguard.py run --wait --cores 16 --mem-gb 64 --label k5-agentic_route_r1 -- /apps/development/rlhf/align/.venv/bin/python /apps/development/rlhf/align/scripts/p2/k5_levers.py route --run-dir /apps/development/rlhf/align/.tmp/p2/k5/dev --round 1 --mode subgoal --workers 2 --bank-k 3 --sub-n 4 --pairs /apps/development/rlhf/align/.tmp/p2/k5/dev/agentic/ver_r0.jsonl /apps/development/rlhf/align/.tmp/p2/k5/dev/agentic/gen_r0.jsonl

# agentic_sg_r1
/apps/development/rlhf/align/.venv/bin/python /apps/development/rlhf/align/scripts/hostguard.py run --wait --cores 24 --mem-gb 100 --label k5-agentic_sg_r1 -- /apps/development/rlhf/align/.venv/bin/python /apps/development/rlhf/align/scripts/p2/pipeline_gen_verify.py --run-dir /apps/development/rlhf/align/.tmp/p2/k5/dev/pipe/agentic_sg_r1 --problems /apps/development/rlhf/align/.tmp/p2/k5/dev/agentic/round1/sg_problems.jsonl --verify-out /apps/development/rlhf/align/.tmp/p2/k5/dev/agentic/round1/sg_ver.jsonl --workers 8 --label k5-agentic_sg_r1 --gen-args '--model /apps/development/rlhf/hf_cache/hub/models--Goedel-LM--Goedel-Prover-V2-8B/snapshots/dfd02e6271a58375dfbf3ece0175277cf6b6a89a --problems /apps/development/rlhf/align/.tmp/p2/k5/dev/agentic/round1/sg_problems.jsonl --out /apps/development/rlhf/align/.tmp/p2/k5/dev/agentic/round1/sg_gen.jsonl --n 4 --cap-tokens 16384 --round 0 --tp 4 --seed 0 --max-num-seqs 128 --gpu-mem-util 0.82 --label k5-agentic_sg_r1 --prompt-format goedel_bank'
```
After the GPU stages: `k5_levers.py report --run-dir $K5 --budget restart_k:8` (also `status`, which lists under-filled problems and how many units each needs; generate those with the matching stage and a larger `--n-restart`/`--n-agentic`, then re-run `report`).

Estimate (`k5_levers.py estimate`; measured rate 522.5 tok/s end-to-end = 1.88 M tokens per box-hour, mean 5,685 tokens/attempt and 36% failure on miniF2F-test from the P2 baseline manifest; ProofNet-valid mean 9,000 tokens and 75% failure are ASSUMED; repair success 12%, repair prompts 1.15x cost, subgoal rounds 0.35 x an attempt per failed parent are ASSUMED):

| Profile | restart | correct | agentic | total (generation only) |
|---|---|---|---|---|
| full: n_restart 32, n_agentic 8 (PREREG pass@8 from n = 32) | 50 | 63 | 23 | **~136 box-h** |
| lean (recommended first pass): n_restart 16, n_agentic 6 | 25 | 31 | 17 | **~74 box-h** |
| minimal: n_restart 8, `--budget restart_k:4`, n_agentic 4 | 13 | 16 | 11 | **~40 box-h** |
| miniF2F-valid only, lean profile | 11 | 9 | 6 | **~26 box-h** |

These exceed the plan's ~30 box-h for K5 except for the miniF2F-valid-only and minimal profiles; the drivers are repairing EVERY failed root (needed for an unbiased replay) and ProofNet's long attempts. Options to cut cost, each a decision for the plan owner: miniF2F-valid first (26 box-h), then ProofNet only if informative; repair only the first 8 roots; or a stratified dev subset. Verification CPU time (6 s/proof order of magnitude at 8 workers, timeouts dominate) overlaps generation in `pipeline_gen_verify.py` for the GPU stages; the CPU-only stages (route, asm verify) add roughly 1-2 h of wall time per agentic round at 2-8 Lean workers.

## 8. Tests

```
cd /apps/development/rlhf/align-wt/agentic   # (or the merged checkout)
. /apps/development/rlhf/env/apps_only.sh; export PYTHONPATH=$PWD
PY=/apps/development/rlhf/align/.venv/bin/python
$PY scripts/hostguard.py run --cores 4 --mem-gb 16 --label k5-tests --wait -- $PY -m pytest tests/test_k5_levers.py tests/test_subgoal_repair.py -q            # CPU, ~5 s (slow tests deselected with -m "not slow")
$PY scripts/hostguard.py run --cores 4 --mem-gb 16 --label k5-lean --wait -- $PY -m pytest tests/test_subgoal_repair.py -q -m slow     # real Lean, ~45 s
```