From Search to Policy: Thompson Sampling Trees for Robust LLM Code Refinement

Iterative refinement has become an effective strategy for improving LLM performance on difficult coding tasks, but existing methods suffer from limited exploration during debugging. We introduce Tree-Guided Policy Refinement (TGPR), which augments GRPO with Thompson Sampling-guided tree search to generate diverse refinement trajectories at training time only, with no additional inference overhead. Across HumanEval, MBPP, APPS, Codeforces, and LiveCodeBench, TGPR consistently outperforms competitive RL baselines, improving pass@1 by up to +4.4 pp and pass@10 by up to +5.2 pp over GRPO, with especially strong gains on harder and contamination-resistant benchmarks. Under zero-shot greedy decoding, TGPR (77.6 MBPP, 41.2 LiveCodeBench pass@1) exceeds GPT-4o-mini (74.8 and 36.7), though the comparison is not fully controlled as TGPR is fine-tuned on task-specific data while GPT-4o-mini is not. These results suggest that structured training-time exploration is a scalable alternative to increasing model scale for iterative code refinement. Code: https://github.com/AnonimusCoders/TGPRcoderefinement.