--- license: apache-2.0 tags: - reinforcement-learning - theorem-proving - lean4 - aesop --- # satp-policy-v2 A 0.2B ByT5 policy that reads a Lean 4 goal state and emits a per-problem `aesop` configuration — search budget, rule-set options, and extra tactic and lemma rules. This repository is a self-contained inference bundle: checkpoint plus a standalone `infer.py` (no training-repo dependencies), forward-equivalent to the training policy. ## Results miniF2F ([`ChristianZ97/minif2f-satp`](https://huggingface.co/datasets/ChristianZ97/minif2f-satp)), greedy decode, statements submitted raw (Lean default `maxHeartbeats 200000`), Lean 4 v4.26.0: | split | solved | |---|---| | test | **37.7% (92/244)** | | validation | 37.3% (91/244) | | ablation (test) | solved | |---|---| | `satp` | **37.7% (92/244)** | | w/ default search budget | 37.7% (92/244) | | distilled static configuration | 31.6% (77/244) | | w/o additional lemma | 31.1% (76/244) | | w/o additional tactic | 30.7% (75/244) | | `aesop` | 10.7% (26/244) | Every number above — both splits and all ablation rows — is verified per problem by two independent backends (the Kimina server, and one `lake env lean` per proof), with `#print axioms` as the arbiter: a counted proof compiles, contains no `sorry`, and depends on no axioms beyond `propext`, `Classical.choice`, `Quot.sound`. Policy rows are additionally checked through both input paths (dataset `goal_state`, and live goal states via the `satp` tactic) with identical solved sets. The 92 solved test proofs are published in two forms — the policy-emitted configuration verbatim, and a distilled minimal tactic script — at [LeanSATP `proofs/minif2f_test_v2`](https://github.com/ChristianZ97/LeanSATP/tree/satp-goal-state-input/proofs/minif2f_test_v2); every file compiles on stock Mathlib v4.26 with `import Mathlib` alone. ## Checkpoint | | | |---|---| | run | `e0f230ns` (wandb `satp-v2`), epoch 7 / step 1700, seed 1827 | | verifier | Kimina-Lean Server, Lean 4 v4.26.0 | | architecture | ByT5(`kaiyuy/leandojo-lean4-retriever-byt5-small`)+LoRA(r16,α32,q/k/v/o) mean-pool → shared MLP → 69-head joint action (28 tactic ×10-way + 32 lemma ×181-way + 4 config_level + 5 config_binary) + dense premise retrieval | ## Distilled static configuration One constant config for every problem (lemma off); per-head mode over the 244 val decodes, `-- %` = mode agreement (joint 10-way: off | type×priority). ```lean aesop (config := { -- maxRuleApplications 200 = default -- 50% maxRuleApplicationDepth := 10 -- 50% maxNormIterations := 520 -- 36% maxGoals := 1024 -- 24% }) (rule_sets := [-builtin]) -- enableBuiltin=F 59% (add norm 100 (by ring)) -- 86% (add norm 100 (by field_simp)) -- 82% (add norm 100 (by norm_num)) -- 93% (add norm 100 (by norm_cast)) -- 82% (add norm 0 (by linarith)) -- 64% (add norm 0 (by nlinarith)) -- 89% (add norm 0 (by omega)) -- 82% (add norm 100 (by abel)) -- 85% (add norm 100 (by zify)) -- 64% (add norm 0 (by bound)) -- 55% (add norm 100 (by interval_cases)) -- 65% (add norm 100 (by ext)) -- 55% (add norm 100 (by split)) -- 60% (add norm 100 (by exfalso)) -- 84% (add norm 100 (by ring_nf at *)) -- 58% (add norm 100 (by field_simp [*] at *)) -- 45% (add norm 100 (by norm_num [*] at *)) -- 72% (add norm 100 (by norm_cast at *)) -- 82% (add norm 100 (by rfl)) -- 86% (add norm 100 (by push_cast)) -- 93% (add norm 100 (by assumption_mod_cast)) -- 72% (add safe 0 (by ring_nf)) -- 32% (add safe 0 (by simp_all)) -- 43% ``` off: positivity 87%, push_neg 85%, gcongr 65%, simp 58%, decide 36%. binary modes: enableSimp T 94%, useSimpAll T 71%, enableUnfold T 63%, useDefaultSimpSet T 86%, enableBuiltin F 59%. Priorities are integers, smaller = higher (0 first, 100 last); unsafe = success %. ## Files ``` best_checkpoint.pt # 1.12 GB — infer.py reads model_state_dict only infer.py # load → retrieve → greedy-decode → Kimina verify reproduce.py # test-split repro, kimina | lake verifier (imports infer.py) cache/premise_embeddings.npy # 1.0 GB — 180,957 × 1472 premise embeddings cache/mathlib4_premises.txt # 38 MB — premise names, index-aligned ``` Premise corpus = LeanDojo Benchmark-4 v10; embeddings come from the frozen retriever → verifier-version independent. ByT5 base auto-downloads on first run. ## Run ```bash pip install torch transformers numpy requests huggingface_hub datasets hf download ChristianZ97/satp-policy-v2 --local-dir satp-policy-v2 cd satp-policy-v2 KIMINA_URL=http://localhost:8000 python infer.py ``` | var | default | meaning | |---|---|---| | `KIMINA_URL` | `http://localhost:8000` | Kimina `/verify` endpoint | | `SATP_SPLIT` | `validation` | or `test` | | `SATP_LIMIT` | `0` | >0 → first N problems only | | `SATP_CKPT` | `./best_checkpoint.pt` | checkpoint path | | `SATP_CACHE` | `./cache` | premise files dir | | `SATP_HEARTBEATS0` | unset | `1` → remove the heartbeat cap (diagnostic; reported numbers use the default) | | `SATP_LEAN_TIMEOUT` | `300` | server-side Lean timeout (s); HTTP timeout = +60 | Verification is sequential (one POST per problem, retried with backoff). Header-less statements get the training header prepended. ## Reproduce `reproduce.py` re-runs the miniF2F **test** split end-to-end over the same greedy decode and compares against 92/244, with either verifier: ```bash KIMINA_URL=http://localhost:8000 python reproduce.py kimina SATP_LAKE_DIR=/path/to/lean-v4.26-project python reproduce.py lake # no server ``` `lake` runs one `lake env lean` per problem (300 s timeout) inside any Lean v4.26 project with Mathlib built; writes `reproduce_ .jsonl` (`{name, success, tactic, code}` per problem). Reference: kimina **92/244** exact (twice, bit-identical — greedy decode is deterministic); all 92 replay **92/92** under `lake env lean`. ### Lean-only path [LeanSATP](https://github.com/ChristianZ97/LeanSATP/tree/satp-goal-state-input) runs the same policy as a Lean tactic: take a dataset row's `formal_statement`, add `import LeanSATP`, and close the proof with `satp` (`satp?` also prints the emitted configuration). Lean handles the rest — elaboration, live goal state → policy, `aesop` execution. Same 92/244 solved set.
