EPIG-Tree: Getting More Credit from Every Rollout

Why Not GRPO? By observing our historical rollouts, we can make better decisions about where to allocate resources in order to learn more effectively.

First, we can simplify the problem into estimating the policy gradient using reward observations from sampled actions. GRPO does this by treating an entire stochastic rollout as a single sample and then comparing its return against the returns of other rollouts in the group, yielding a per-rollout advantage estimate. Where this fails is when rollouts become sparse, for example when the model executes many concurrent actions. Intuitively, this becomes an issue because the signal in what is good or bad becomes weaker relative to the noise from factors that are irrelevant to the reward.

What EPIG-Tree proposes is to decompose the advantage of an entire rollout into its constituent parts, allowing us to obtain a better estimate for each part. This has previously been explored in approaches such as TreeRL, where per-token entropy is used to determine branching points in the rollout tree. In other words, when the model reaches a token it is uncertain about, a branch is created.

Intuitively, this makes sense because it allows us to explore multiple distinct but likely paths. We can view training as the process of reducing the model's uncertainty, so branching at points of uncertainty gives the training algorithm an opportunity to make the model more certain by revealing which paths lead to better outcomes.

TreeRL comes with two major flaws. First, entropy does not necessarily approximate the model's uncertainty with respect to the task. Modern LLMs can exhibit high entropy during reasoning or formulation, even when the choice being made is orthogonal to the task itself. Second, TreeRL assumes that we have no prior information about the task. However, if we reframe training as a sequential process, we can leverage previous rollouts to optimize where additional branching would be most informative.

With this framing, the problem can now be formalized as an optimization problem using observation.

The algorithm, in plain terms

1. Run a few ordinary rollouts and record the states, actions, log-probs, rewards, and costs.
2. From that history, propose candidate branch points and estimate three things at each: how different the options' values look, how much the decision moves the model, and how expensive a rollout is.
3. Score the candidates with EPIG, branch at the top ones, and re-roll the noisy children according to the suffix law.
4. Turn the local values into advantages and apply a normal clipped update.

Results

We checked this in three layers: a frozen setting where the true gradient is (almost) computable, frozen language-model diagnostics, and online training.

Cloned-state control. Freeze the model, clone a state, and compare each budget-limited estimate against a high-budget reference gradient. EPIG beats both entropy and uniform on gradient error in 9 of 13 environments - exactly the nine dense, smooth-reward ones - and on those it recovers the gradient direction almost perfectly (cosine 0.998–1.000). The four misses are the tiny-action and sparse-reward cases, where there is nothing to allocate.

Frozen LLM gradients. On frozen GSM8K with Qwen3-8B, EPIG's tree gradient is both better aligned and far steadier than entropy branching (cosine 0.075 vs 0.022). Flat GRPO aligns best of all - a reminder that branch placement and the local advantage estimator are separate questions.

Single-turn math. Tree-local credit beats flat GRPO, but the gaps between placement methods sit within single-seed noise - single-turn math is too saturated to isolate placement.

Frozen Wordle. EPIG reaches the lowest value error (0.523 vs 0.568 entropy, 1.094 uniform). The telling bar is the last one: branching on plain value spread, without the score weighting, is much worse (1.389) - the score term is doing real work.

Online Wordle. Trained for 300 updates, every tree method starts slower than flat GRPO, but flat GRPO saturates near 0.790 and stops. The tree methods keep climbing and overtake it; EPIG finishes highest at 0.850 (vs 0.825 entropy, 0.805 uniform, 0.790 flat), gaining the most over the second half - selective placement compounds once the policy reaches states where the options genuinely differ.

Paper

The full derivation - the variance decomposition, the branch-placement and suffix-allocation laws, and all experiments - is in our paper, EPIG-Tree: Compute-Optimal Branching for Gradient-Efficient Reinforcement Learning.