Learning State Transitions Improves Downstream Policy Learning in Geometry Theorem Proving

Geometry theorem proving can be formulated as sequential action selection over an evolving proof state. Standard policy supervised fine-tuning (Policy-SFT) teaches a model to imitate expert actions, but does not directly supervise how those actions change the state. We study whether explicitly learning these state transitions before policy training improves downstream theorem proving. We build a stateful agent environment around the open-source GenesisGeo symbolic engine. Within this environment, we first apply world-model supervised fine-tuning (WM-SFT) to predict construction feedback, whose newly derived geometric relations define the action-induced state delta. We then perform Policy-SFT on expert actions using the same model. WM-SFT followed by Policy-SFT achieves 74.0\% proof success, compared with 63.2\% for Policy-SFT alone. Additionally, after subsequent reinforcement learning, WM-SFT followed by Policy-SFT and RL reaches up to 76.8\%, compared with 71.2\% for Policy-SFT followed by RL. WM-SFT is used only during training; all policies follow the same action-selection procedure at inference, with no additional prediction stage, model call, or planning. These results show that supervised learning of action-induced state transitions can provide an effective initialization for training mathematical-agent policies.