EPIG-Tree: Compute-Optimal Branching for Gradient-Efficient Reinforcement Learning
Khomich, Hermansson and Hakimi treat tree-structured rollout construction in LLM RL as a compute-allocation problem, deriving where to branch from a law-of-total-variance decomposition of the local policy gradient rather than from policy entropy.
Ask this paper
GRPO collapses a trajectory to one scalar. Every token inherits a single trajectory-level advantage even though the trajectory contains many distinct causal decisions, recovery attempts and environment-randomness events.
Branch where it reduces gradient uncertainty, not where the policy is uncertain. The decomposition yields two allocation laws: new branches reduce decision uncertainty, repeated suffix rollouts reduce continuation uncertainty, and the EPIG-Tree score combines both using rollouts already computed.
The suffix law is explicit, allocating rollouts proportional to occupancy weight times score-function norm times value standard deviation, divided by the square root of cost.
Gradient MSE falls across a controlled sweep. EPIG wins in all nine dense continuous-control environments of a 13-environment sweep and recovers the reference gradient direction nearly exactly, and it improves frozen-LLM gradient calibration over entropy branching.
Branch placement matters less than token-level credit in single-turn math, but in multi-turn Wordle EPIG reaches a 0.850 final win rate against flat GRPO's early saturation at 0.790, so the advantage appears in stateful settings with large action spaces.
Abstract
Reward-based reinforcement learning for language models, exemplified by Group Relative Policy Optimization (GRPO), collapses an entire stochastic trajectory into a single scalar reward. This is clean and scalable, but it explores and allocates reward inefficiently: a trajectory may contain many causal decisions, recovery attempts, and environment-randomness events, yet every token or action inherits one trajectory-level advantage. We study tree-based rollout construction as a compute-allocation problem for policy-gradient estimation. Our central claim is that branches should be placed not where the policy is merely uncertain, but where an additional branch most reduces uncertainty about the policy gradient per unit of compute. From a law-of-total-variance decomposition of the local policy-gradient random variable, we derive two allocation laws: new branches reduce decision uncertainty, while repeated suffix rollouts reduce continuation uncertainty. The resulting EPIG-Tree score allocates branches using the already computed rollouts. It estimates occupancy- and score-weighted value uncertainty, along with a suffix law $n_e \propto w_e \|\nabla_θ\log π(a_e|h_e)\| σ_e / \sqrt{c_e}$. Empirically, EPIG reduces gradient MSE in cloned-state control, winning in all nine dense continuous-control environments of a 13-environment sweep and recovering the reference gradient direction near-perfectly, and it improves frozen-LLM gradient calibration relative to entropy branching. In online single-turn math, tree-local credit beats flat GRPO, while branch placement is secondary to token-level credit assignment. In online multi-turn Wordle, EPIG attains the highest final win rate (0.850), overtaking flat GRPO, which saturates early at 0.790, and entropy branching as training proceeds, confirming that the gradient-estimation advantage transfers to a stateful, large-action setting.