EPIG-Tree: Compute-Optimal Branching for Gradient-Efficient Reinforcement Learning

TL;DR - EPIG-Tree is a compute-aware branching method for reinforcement learning that places branches where they most reduce policy-gradient uncertainty. It improves gradient estimation and outperforms flat GRPO and entropy-based branching in several control and language-model tasks.

• EPIG-Tree distinguishes decision uncertainty, addressed with new branches, from continuation uncertainty, addressed with repeated suffix rollouts.
• Its allocation score uses existing rollouts to estimate occupancy- and gradient-weighted value uncertainty while accounting for rollout cost.
• EPIG reduced gradient mean-squared error across nine dense continuous-control environments and improved frozen-LLM gradient calibration.
• In multi-turn Wordle, EPIG reached a 0.850 final win rate versus 0.790 for flat GRPO; in single-turn math, token-level credit assignment mattered more than branch placement.