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

# EPIG-Tree: Compute-Optimal Branching for Gradient-Efficient Reinforcement Learning
Source: [https://arxiv.org/html/2609.20004](https://arxiv.org/html/2609.20004)
###### 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 tree branches should not be placed where the policy is merely uncertain; they should be placed where an additional branch most reduces uncertainty about the policy gradient accounting for the expected compute cost\. 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

ne∝we​‖∇θ​log​π​\(ae∣he\)‖​σece\.n\_\{e\}\\propto\\frac\{w\_\{e\}\\,\\\|\\nabla\_\{\\theta\}\\log\\pi\(a\_\{e\}\\mid h\_\{e\}\)\\\|\\,\\sigma\_\{e\}\}\{\\sqrt\{c\_\{e\}\}\}\.Empirically, we show that EPIG reduces gradient MSE in cloned\-state control, winning in all nine dense continuous\-control environments of a thirteen\-environment sweep and recovering the reference gradient direction near\-perfectly, and it improves frozen\-LLM gradient calibration relative to simpler 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\.

## 1Introduction

The recent success of reinforcement learning with verifiable rewards has made on\-policy algorithms such as GRPO central to LLM reasoning training\[[1](https://arxiv.org/html/2609.20004#bib.bib1)\]\. GRPO samples a group of responses for the same prompt, normalizes terminal rewards within the group, and applies a PPO\-style clipped objective without training a separate critic\.

GRPO nonetheless has three structural failure modes in long\-horizon or multi\-turn settings\.

Poor recovery credit assignment\.Suppose that a trajectory makes a mistake while performing correct actions and fails\. Reward based GRPO assigns the whole trajectory a negative group\-relative reward, so the correct actions are also punished\. Conversely, if a trajectory succeeds despite multiple bad actions, the bad actions are reinforced\. Recovery behavior is thus not a special case of long\-horizon credit assignment; it is a direct consequence of assigning one scalar to a whole path\.

Inefficient handling of nondeterminism\.In a stochastic environment, a terminal reward is a noisy observation ofQ\(h,a\)=𝔼\[R∣h,a\]Q\(h,a\)=\\mathbb\{E\}\[R\\mid h,a\]\. A single trajectory conflates policy choice with environment randomness\. If the environment randomly sabotages a good action, GRPO punishes the action; if it rescues a bad action, GRPO reinforces it\. Estimating conditional continuation values requires repeated counterfactual suffixes from the same state\.

Poor scaling with rollout count\.Increasing the number of independent full rollouts explores more, yet many rollouts differ only in low\-importance tokens, formatting, or environment noise\. These do not provide any signal for the model to learn with\.

Tree rollouts are a natural remedy\. Instead of samplingKKindependent full trajectories, we reuse prefixes, branch at intermediate states, and estimate descendant values\. This yields local, process\-like advantages without training a separate process reward model\. Recent LLM tree methods follow this direction: TreeRL’s EPTree forks from high\-entropy intermediate tokens\[[2](https://arxiv.org/html/2609.20004#bib.bib2)\]; TreePO uses local uncertainty and segment\-level tree modeling to amortize common prefixes\[[3](https://arxiv.org/html/2609.20004#bib.bib3)\]; and related tree\-structured rollout methods derive relative advantages in multi\-turn agent tasks\. The unresolved question is:*where should the tree branch?*

A common answer is entropy\. Entropy is useful: if a policy is deterministic at a prefix, its branches coincide\. But entropy alone is not the quantity that the policy\-gradient estimator cares about\. A high\-entropy choice among reward\-equivalent phrasings is not worth branching; a moderate\-entropy choice between high\- and low\-value actions can be crucial\. The organizing principle of this paper is therefore:

> A tree is an experimental apparatus for estimating the policy gradient\. Optimal branching allocates compute where an additional branch most reduces gradient\-estimator error per unit cost—not where the policy is merely uncertain\.

Concretely: GRPO collapses a full trajectory into one scalar advantage, whereas a tree estimates local conditional values; entropy measures only whether branches can differ, while EPIG measures whether they matter for the gradient\.

Contributions\.

1. 1\.We derive a law\-of\-total\-variance decomposition showing that tree construction has two distinct jobs: new branches reduce decision uncertainty, while suffix resampling reduces continuation uncertainty\.
2. 2\.We prove a compute\-optimal suffix allocation law ne∗∝we​‖∇log⁡π​\(ae∣he\)‖​σecen\_\{e\}^\{\*\}\\propto\\frac\{w\_\{e\}\\,\\\|\\nabla\\log\\pi\(a\_\{e\}\\mid h\_\{e\}\)\\\|\\,\\sigma\_\{e\}\}\{\\sqrt\{c\_\{e\}\}\}and a marginal branching law based on occupancy and score weighted value uncertainty\.
3. 3\.We introduce EPIG\-Tree, an algorithm that uses entropy as a proposal mechanism but allocates branches by expected predictive information gain about the gradient\.
4. 4\.We validate the theory across cloned\-state continuous control, frozen LLM gradient calibration, single\-turn LLM math, and online multi\-turn Wordle, establishing both where EPIG wins and the boundary conditions under which it ties cheaper baselines\.

## 2Related work

Policy gradients and PPO\.Our derivation begins from the standard policy\-gradient theorem\[[4](https://arxiv.org/html/2609.20004#bib.bib4)\]and PPO\-style clipped optimization\[[5](https://arxiv.org/html/2609.20004#bib.bib5)\]\. The novelty is not a new policy\-gradient identity; it is the use of the identity to formulate tree construction as optimal sampling for the gradient\.

GRPO and RLVR\.DeepSeekMath introduced GRPO as a critic\-free PPO variant for mathematical reasoning\[[1](https://arxiv.org/html/2609.20004#bib.bib1)\], replacing a learned value baseline with group\-relative normalization\. Our critique is that the group\-normalized terminal reward is a coarse estimator of local advantages in long\-horizon settings\.

Tree search for LLM RL\.TreeRL proposes EPTree, an entropy\-guided tree search that forks from high\-uncertainty tokens and derives process supervision from descendant correctness\[[2](https://arxiv.org/html/2609.20004#bib.bib2)\]\. TreePO similarly uses local uncertainty with segment\-level tree modeling and prefix amortization\[[3](https://arxiv.org/html/2609.20004#bib.bib3)\]\. These methods motivate our setting; we replace entropy/uncertainty heuristics with a gradient\-estimation objective\.

Entropy in LLM RL\.The entropy\-mechanism analysis argues that RL for reasoning models trades policy entropy for downstream performance, deriving entropy change from the covariance between logit updates and action probabilities/advantages\[[6](https://arxiv.org/html/2609.20004#bib.bib6)\]\. Complementary work finds that a minority of high\-entropy tokens act as reasoning forks and carry much of the RLVR update signal\[[7](https://arxiv.org/html/2609.20004#bib.bib7)\]\. We agree that entropy identifies branchable or update\-sensitive points\. Our result is narrower and sharper: entropy is not the objective; it is one input to a compute\-allocation problem\.

Reward overoptimization\.Gao, Schulman, and Hilton study gold reward as a function of optimization distanced=DKL\(π∥π0\)d=\\sqrt\{D\_\{\\mathrm\{KL\}\}\(\\pi\\\|\\pi\_\{0\}\)\}, finding distinct scaling forms for best\-of\-nnand RL optimization and Goodhart\-style degradation under proxy rewards\[[8](https://arxiv.org/html/2609.20004#bib.bib8)\]\. We include a Goodhart\-corrected EPIG\-Full variant for proxy\-reward settings, while our strongest evidence concerns exact or cloned\-state gradient estimation\.

## 3From outcome GRPO to local gradient experiments

Consider a history or prefixhh, an actiona∼πθ\(⋅∣h\)a\\sim\\pi\_\{\\theta\}\(\\cdot\\mid h\), and a utility

Y=U⁡\(G⁡\(h\)\+Rfuture\),Y=U\\bigl\(G\(h\)\+R\_\{\\mathrm\{future\}\}\\bigr\),\(1\)whereG⁡\(h\)G\(h\)is reward already accumulated andUUis the training utility \(linear return, success/failure, preference utility, verifier score, or task\-specific reward\)\. Define

Qh\(a\)=𝔼\[Y∣h,a\],σh2\(a\)=Var\(Y∣h,a\),ψh\(a\)=∇θlogπθ\(a∣h\)\.Q\_\{h\}\(a\)=\\mathbb\{E\}\[Y\\mid h,a\],\\qquad\\sigma\_\{h\}^\{2\}\(a\)=\\mathrm\{Var\}\(Y\\mid h,a\),\\qquad\\psi\_\{h\}\(a\)=\\nabla\_\{\\theta\}\\log\\pi\_\{\\theta\}\(a\\mid h\)\.\(2\)The local contribution to the policy gradient is

gh=μh​𝔼a∼πh​\[ψh​\(a\)​Qh​\(a\)\],g\_\{h\}=\\mu\_\{h\}\\,\\mathbb\{E\}\_\{a\\sim\\pi\_\{h\}\}\\bigl\[\\psi\_\{h\}\(a\)Q\_\{h\}\(a\)\\bigr\],\(3\)whereμh\\mu\_\{h\}is the occupancy or training weight of the prefix\. For an LLM token prefix,μh\\mu\_\{h\}captures how frequently the current policy reaches that prefix or how much training mass we assign to it; for a cloned Gym/MuJoCo state,μh\\mu\_\{h\}may be uniform over sampled states\.

Outcome\-only GRPO estimates many local terms with a single trajectory scalar\. WithKKsampled completionsτi\\tau\_\{i\}for a prompt, the group advantage is

A^iGRPO=Ri−R¯sR\+ε,\\widehat\{A\}\_\{i\}^\{\\mathrm\{GRPO\}\}=\\frac\{R\_\{i\}\-\\bar\{R\}\}\{s\_\{R\}\+\\varepsilon\},\(4\)and every action on trajectoryiireceives this same sign and scale\. A tree estimator instead estimatesQh​\(a\)Q\_\{h\}\(a\)by descendant rollouts from a shared prefix:

Q^h​\(a\)=1nh,a​∑j=1nh,aYh,a,j,A^​\(h,a\)=Q^h​\(a\)−V^​\(h\)\.\\widehat\{Q\}\_\{h\}\(a\)=\\frac\{1\}\{n\_\{h,a\}\}\\sum\_\{j=1\}^\{n\_\{h,a\}\}Y\_\{h,a,j\},\\qquad\\widehat\{A\}\(h,a\)=\\widehat\{Q\}\_\{h\}\(a\)\-\\widehat\{V\}\(h\)\.\(5\)The mathematical question is now: given a compute budget, where should we spend the next suffix rollout or branch?

## 4The variance decomposition that determines the tree

Let the random vector for a one\-sample local gradient contribution be

Zh=μh​ψh​\(A\)​Y,A∼πh\.Z\_\{h\}=\\mu\_\{h\}\\psi\_\{h\}\(A\)Y,\\qquad A\\sim\\pi\_\{h\}\.\(6\)Its total variance decomposes as follows\.

Proposition 1 \(Decision and continuation uncertainty\)\.For fixed prefixhh,

Var⁡\(Zh∣h\)\\displaystyle\\mathrm\{Var\}\(Z\_\{h\}\\mid h\)=μh2​VarA∼πh​\[ψh​\(A\)​Qh​\(A\)\]\\displaystyle=\\mu\_\{h\}^\{2\}\\,\\mathrm\{Var\}\_\{A\\sim\\pi\_\{h\}\}\\bigl\[\\psi\_\{h\}\(A\)Q\_\{h\}\(A\)\\bigr\]\(7\)\+μh2​𝔼A∼πh​\[ψh​\(A\)​ψh​\(A\)⊤​σh2​\(A\)\]\.\\displaystyle\\quad\+\\mu\_\{h\}^\{2\}\\,\\mathbb\{E\}\_\{A\\sim\\pi\_\{h\}\}\\bigl\[\\psi\_\{h\}\(A\)\\psi\_\{h\}\(A\)^\{\\top\}\\sigma\_\{h\}^\{2\}\(A\)\\bigr\]\.\(8\)Taking the trace gives a scalar gradient\-MSE proxy\.

Proof\.Apply the law of total variance,Var\(Zh∣h\)=Var\(𝔼\[Zh∣A,h\]∣h\)\+𝔼\[Var\(Zh∣A,h\)∣h\]\\mathrm\{Var\}\(Z\_\{h\}\\mid h\)=\\mathrm\{Var\}\(\\mathbb\{E\}\[Z\_\{h\}\\mid A,h\]\\mid h\)\+\\mathbb\{E\}\[\\mathrm\{Var\}\(Z\_\{h\}\\mid A,h\)\\mid h\]\. Since𝔼\[Zh∣A=a,h\]=μhψh\(a\)Qh\(a\)\\mathbb\{E\}\[Z\_\{h\}\\mid A=a,h\]=\\mu\_\{h\}\\psi\_\{h\}\(a\)Q\_\{h\}\(a\)andVar⁡\(Y∣h,a\)=σh2​\(a\)\\mathrm\{Var\}\(Y\\mid h,a\)=\\sigma\_\{h\}^\{2\}\(a\), the result follows\.□\\square

This decomposition is the core of the paper\. The first term is*decision uncertainty*: uncertainty over which action branch matters\. The second is*continuation uncertainty*: noise in the suffix after a branch action has already been selected\. Tree construction therefore has two distinct compute actions:

- •add a new branch fromhhto reduceVara​\[ψh​\(a\)​Qh​\(a\)\]\\mathrm\{Var\}\_\{a\}\[\\psi\_\{h\}\(a\)Q\_\{h\}\(a\)\];
- •add suffix samples under an existing edge\(h,a\)\(h,a\)to reduceσh2​\(a\)\\sigma\_\{h\}^\{2\}\(a\)\.

### 4\.1Why entropy is insufficient

Entropy isHh=H\(πθ\(⋅∣h\)\)H\_\{h\}=H\(\\pi\_\{\\theta\}\(\\cdot\\mid h\)\): it measures how many alternatives the policy can produce\. But the variance decomposition contains no standalone entropy term\. A high\-entropy prefix withQh​\(a\)Q\_\{h\}\(a\)nearly constant has little gradient\-relevant decision uncertainty; a moderate\-entropy prefix with largeQhQ\_\{h\}dispersion can dominate the gradient MSE\. Entropy is therefore a useful proposal mechanism—it tells us where branches can differ\. EPIG asks whether those differences matter\.

## 5Optimal allocation laws

### 5\.1Suffix allocation

Let edgee=\(h,a\)e=\(h,a\)have training weightwew\_\{e\}, score vectorψe=∇log⁡π​\(a∣h\)\\psi\_\{e\}=\\nabla\\log\\pi\(a\\mid h\), continuation standard deviationσe\\sigma\_\{e\}, and suffix costcec\_\{e\}\. Ifnen\_\{e\}independent suffixes are allocated underee, the trace\-variance contribution is approximatelyAe/neA\_\{e\}/n\_\{e\}with

Ae=we2​‖ψe‖2​σe2\.A\_\{e\}=w\_\{e\}^\{2\}\\\|\\psi\_\{e\}\\\|^\{2\}\\sigma\_\{e\}^\{2\}\.
Theorem 1 \(Cost\-sensitive suffix allocation\)\.For fixed candidate edges and budgetBB, the solution of

min⁡∑ene\>0⁡Aenes\.t\.∑ece​ne≤B\\min\_\{n\_\{e\}\>0\}\\sum\_\{e\}\\frac\{A\_\{e\}\}\{n\_\{e\}\}\\qquad\\text\{s\.t\.\}\\qquad\\sum\_\{e\}c\_\{e\}n\_\{e\}\\leq Bis

ne∗∝Aece=we​‖ψe‖​σece\.n\_\{e\}^\{\*\}\\propto\\sqrt\{\\frac\{A\_\{e\}\}\{c\_\{e\}\}\}=\\frac\{w\_\{e\}\\\|\\psi\_\{e\}\\\|\\sigma\_\{e\}\}\{\\sqrt\{c\_\{e\}\}\}\.\(9\)
Proof\.The Lagrangian isℒ⁡\(n,ρ\)=∑eAe/ne\+ρ⁡\(∑ece​ne−B\)\\mathcal\{L\}\(n,\\rho\)=\\sum\_\{e\}A\_\{e\}/n\_\{e\}\+\\rho\(\\sum\_\{e\}c\_\{e\}n\_\{e\}\-B\)\. Stationarity gives−Ae/ne2\+ρce=0\-A\_\{e\}/n\_\{e\}^\{2\}\+\\rho c\_\{e\}=0, hencene=Ae/\(ρ​ce\)n\_\{e\}=\\sqrt\{A\_\{e\}/\(\\rho c\_\{e\}\)\}; the proportionality follows after normalizing to the budget\.□\\square

The law is a Neyman\-allocation analogue for policy gradients: sample more where gradient leverage and suffix noise are high and cost is low\.

### 5\.2Branch allocation

Letmhm\_\{h\}be the number of action branches already sampled at nodehh, and define the decision\-uncertainty coefficient

Bh=μh2​tr​\(Vara∼πh​\[ψh​\(a\)​Qh​\(a\)\]\)\.B\_\{h\}=\\mu\_\{h\}^\{2\}\\,\\mathrm\{tr\}\\\!\\left\(\\mathrm\{Var\}\_\{a\\sim\\pi\_\{h\}\}\[\\psi\_\{h\}\(a\)Q\_\{h\}\(a\)\]\\right\)\.Withmhm\_\{h\}independent branch samples, the action Monte Carlo term scales asBh/mhB\_\{h\}/m\_\{h\}, so the discrete gain from one more branch is

Δhbranch≈Bhmh​\(mh\+1\)−λ​ch,\\Delta\_\{h\}^\{\\mathrm\{branch\}\}\\approx\\frac\{B\_\{h\}\}\{m\_\{h\}\(m\_\{h\}\+1\)\}\-\\lambda c\_\{h\},\(10\)wherechc\_\{h\}is expected branch/suffix cost andλ\\lambdais compute price\. In practiceQhQ\_\{h\}andψh\\psi\_\{h\}are estimated from pilot branches, and EPIG uses the empirical score

SEPIG​\(h\)=μ^h2​tr​Var^a∈𝒞⁡\(h\)​\[ψ^h​\(a\)​Q^h​\(a\)\]c^h​\(mh\+1\)2\+ε\.S\_\{\\mathrm\{EPIG\}\}\(h\)=\\frac\{\\widehat\{\\mu\}\_\{h\}^\{\\,2\}\\,\\mathrm\{tr\}\\,\\widehat\{\\mathrm\{Var\}\}\_\{a\\in\\mathcal\{C\}\(h\)\}\\\!\\left\[\\widehat\{\\psi\}\_\{h\}\(a\)\\widehat\{Q\}\_\{h\}\(a\)\\right\]\}\{\\widehat\{c\}\_\{h\}\(m\_\{h\}\+1\)^\{2\}\+\\varepsilon\}\.\(11\)We call the fullψ\\psi\-weighted score EPIG\-grad—the variant used in all experiments below—and the value\-variance\-only scoreVar^a​\[Q^h​\(a\)\]/c^h\\widehat\{\\mathrm\{Var\}\}\_\{a\}\[\\widehat\{Q\}\_\{h\}\(a\)\]/\\widehat\{c\}\_\{h\}EPIG\-Lite\. EPIG\-Lite is cheaper and suffices in dense continuous control but can fail where score norms differ or action probabilities are very uneven\. Our frozen Wordle calibration confirms theψ\\psi\-weighting is active: EPIG\-grad beats pure value variance\.

Corollary 1 \(Entropy branching as a special case\)\.Entropy\-guided branching coincides with EPIG only under a homogeneity assumption: across candidate nodes,μh\\mu\_\{h\},‖ψh‖\\\|\\psi\_\{h\}\\\|, value dispersion, suffix noise, and cost must be constant or monotone functions of entropy\. When these quantities decouple, entropy selects branchable but low\-information prefixes\.

### 5\.3Goodhart and entropy\-pricing extensions

When rewards are exact or verifiable, we set Goodhart sensitivity to zero\. For proxy rewards, local tree expansion can overoptimize the proxy\. Following the empirical scaling\-law view of reward\-model overoptimization\[[8](https://arxiv.org/html/2609.20004#bib.bib8)\], define local optimization distancedh=DKL\(πhtree∥πh\)d\_\{h\}=\\sqrt\{D\_\{\\mathrm\{KL\}\}\(\\pi\_\{h\}^\{\\mathrm\{tree\}\}\\\|\\pi\_\{h\}\)\}\. A Goodhart\-corrected gain is

Gh​\(d\)=∫0dαh​eHh​\(u\)​𝑑u−βh​Φ​\(d\),G\_\{h\}\(d\)=\\int\_\{0\}^\{d\}\\alpha\_\{h\}e^\{H\_\{h\}\(u\)\}\\,du\-\\beta\_\{h\}\\Phi\(d\),\(12\)withΦ⁡\(d\)=d2\\Phi\(d\)=d^\{2\}for best\-of\-branch\-like selection andΦ⁡\(d\)=d​log⁡d\\Phi\(d\)=d\\log dfor RL\-like updates\. Ifκh​\(d\)=−Hh′​\(d\)\\kappa\_\{h\}\(d\)=\-H\_\{h\}^\{\\prime\}\(d\)is entropy drain, the marginal equilibrium condition is

Bhmh2\+μh​\[αh​eHh​\(dh\)−βh​Φ′​\(dh\)\]​dh′​\(mh\)=λ​Ch′​\(mh\)\+τ​κh​\(dh\)​dh′​\(mh\)\.\\frac\{B\_\{h\}\}\{m\_\{h\}^\{2\}\}\+\\mu\_\{h\}\\left\[\\alpha\_\{h\}e^\{H\_\{h\}\(d\_\{h\}\)\}\-\\beta\_\{h\}\\Phi^\{\\prime\}\(d\_\{h\}\)\\right\]d\_\{h\}^\{\\prime\}\(m\_\{h\}\)=\\lambda C\_\{h\}^\{\\prime\}\(m\_\{h\}\)\+\\tau\\kappa\_\{h\}\(d\_\{h\}\)d\_\{h\}^\{\\prime\}\(m\_\{h\}\)\.\(13\)We treat this as EPIG\-Full\. Our strongest measured evidence concerns the first term, the branch score, and the suffix law; the Goodhart and entropy\-drain terms are theoretically motivated and remain to be stress\-tested at larger proxy\-reward scale\.

## 6Algorithm

EPIG\-Tree uses entropy to propose candidate branch points but value/gradient information to allocate the branch budget\.

Algorithm 1EPIG\-Tree update for one prompt/state batch0:policy

πθ\\pi\_\{\\theta\}; prompts/states; root count

K0K\_\{0\}; branch budget

BB; pilot size

m0m\_\{0\}; compute price

λ\\lambda
1:Sample

K0K\_\{0\}on\-policy root trajectories; store prefixes/states

hh, actions

aa, log\-probs, rewards, and costs\.

2:Propose a candidate set

𝒞\\mathcal\{C\}of branch states: turn or reasoning\-step boundaries / high\-entropy token segments \(LLMs\), or sampled clone states \(cloneable RL\)\.

3:foreach candidate node

h∈𝒞h\\in\\mathcal\{C\}do

4:Score

SEPIG​\(h\)S\_\{\\mathrm\{EPIG\}\}\(h\)using the branch score above\.

5:endfor

6:whilebranch budget

BBremainsdo

7:Add a branch at

arg⁡maxh​SEPIG​\(h\)\\arg\\max\_\{h\}S\_\{\\mathrm\{EPIG\}\}\(h\)with positive marginal gain; refresh its score\.

8:endwhile

9:Allocate suffix rollouts per edge by

ne∗∝we​‖ψe‖​σe/cen\_\{e\}^\{\*\}\\propto w\_\{e\}\\\|\\psi\_\{e\}\\\|\\sigma\_\{e\}/\\sqrt\{c\_\{e\}\}\.

10:Form tree values

V^​\(h\)\\widehat\{V\}\(h\)from descendant leaves and edge advantages

A^​\(h,a\)=V^​\(child\)−V^​\(h\)\\widehat\{A\}\(h,a\)=\\widehat\{V\}\(\\mathrm\{child\}\)\-\\widehat\{V\}\(h\)\(optionally mixed with a root\-relative advantage\)\.

11:Apply a PPO/GRPO\-style clipped update with an action\-token / branch\-segment mask, so the advantage trains the tokens that caused the branch\.

12:returnupdated policy

πθ\\pi\_\{\\theta\}\.

## 7Experiments

We organize the evidence into three layers\. The first directly tests the allocation law under frozen policies and cloneable states; the second tests frozen LLM gradient calibration; the third tests online LLM RL, where optimization dynamics and token credit assignment interact with branch placement\. All numbers are measured\.

### 7\.1Cloned\-state control: direct validation of the allocation law

Setup\.We freeze a policy, sample states, clone each state, sample candidate actions, and estimate a high\-budget reference gradient with 32 suffix rollouts per action\. Each allocation method receives budgetB=8B=8suffixes and estimatesg^\\widehat\{g\}; we report‖g^−gref‖2\\\|\\widehat\{g\}\-g\_\{\\mathrm\{ref\}\}\\\|\_\{2\}and cosine to the reference\. We compare EPIG\-grad against entropy, uniform, value variance, suffix\-⋅\\sqrt\{\\cdot\}, and oracle variants; the key comparison is EPIG\-grad versus entropy and uniform\.

Result\.EPIG\-grad beats both entropy and uniform on gradient\-MSE in 9/13 environments—exactly the nine dense continuous\-control environments \(Figure[1](https://arxiv.org/html/2609.20004#S7.F1)\)\. On these nine it recovers the reference gradient direction near\-perfectly \(cosine 0\.998–1\.000\), while entropy lags and mis\-points \(cosine 0\.79–0\.94; Figure[2](https://arxiv.org/html/2609.20004#S7.F2)\)\. The four exceptions are informative rather than arbitrary\. Tiny\-action environments \(Acrobot withK=3K=3, a stochastic CartPole withK=2K=2\) leave little allocation problem—uniform already covers the relevant actions—and sparse or non\-smooth reward regimes \(MountainCarContinuous, Pendulum\) make small\-pilot value dispersion unreliable\. EPIG is therefore not a universal replacement for exploration; it is a compute\-allocation law for settings where local value differences can be estimated\.

![Refer to caption](https://arxiv.org/html/2609.20004v1/figures/fig1.png)Figure 1:Cloned\-state control: entropy and uniform gradient\-MSE divided by EPIG MSE \(log scale\)\. Values above 1 mean EPIG is better\. EPIG wins on all nine dense continuous\-control tasks and on the four sparse or tiny\-action exceptions falls back toward the cheaper baselines\.![Refer to caption](https://arxiv.org/html/2609.20004v1/figures/fig2.png)Figure 2:Gradient cosine to the high\-budget reference\. On the nine dense continuous\-control tasks EPIG recovers the reference direction near\-perfectly \(0\.998–1\.000\) while entropy mis\-points\.
### 7\.2Counterfactual branching helps when there is headroom

Before selective EPIG allocation, a simpler question is whether counterfactual branching helps at all\. In a seven\-environment ablation, uniform per\-edge stochastic branching—each edge forks with probabilitypp, a fresh action is resampled, and all root and branch samples feed one group\-normalized GRPO update—improves early learning wherever learning happens \(Figure[3](https://arxiv.org/html/2609.20004#S7.F3)\)\. Relative to flat GRPO \(p=0p=0\) under a capped update budget, InvertedDoublePendulum gains \+52 final return, Hopper \+16, and Reacher \+5 atp=0\.5p=0\.5, with the bulk of the gain captured by a moderatep≈0\.1p\\approx 0\.1–0\.20\.2\. The mechanism is simply more counterfactual signal per update; the cost is proportionally more compute \(branched\-edge fraction→0\.9\+\\to 0\.9\+at highpp\), and gains diminish beyond the sweet spot\. EPIG asks how to spend this branching budget more selectively\.

![Refer to caption](https://arxiv.org/html/2609.20004v1/figures/fig3.png)Figure 3:Even unoptimized stochastic branching supplies useful counterfactual signal in control\. Left: normalized final\-return gain over flat GRPO versus branch probabilityppfor the three environments with headroom; the shaded band marks thep≈0\.1p\\approx 0\.1–0\.20\.2sweet spot\. Right: the branched\-edge fraction \(compute cost\) rises steeply withpp, so highppbuys little extra return\.
### 7\.3Frozen LLM gradient calibration

Setup\.We use Qwen3\-8B on GSM8K and compare budgeted\-tree PPO gradients against a 64\-leaf high\-budget reference over 64 prompts and 3 seeds\. This is not an online training result; it asks whether the branch topology improves the gradient estimator\.

Result\.Among tree methods, EPIG\-Lite is more stable and better aligned than entropy branching: cosine0\.0748±0\.01190\.0748\\pm 0\.0119versus0\.0217±0\.03890\.0217\\pm 0\.0389\(entropy goes negative on one seed\), a gap of \+0\.053 at one\-third the standard deviation \(Figure[4](https://arxiv.org/html/2609.20004#S7.F4)\)\. Flat GRPO aligns best with the GRPO\-style reference \(0\.2298\), which we read as a warning that tree\-local advantages and token masks are themselves consequential estimator choices—branch scoring is not the only thing a tree estimator must get right\.

![Refer to caption](https://arxiv.org/html/2609.20004v1/figures/fig4.png)Figure 4:Frozen GSM8K/Qwen3\-8B gradient calibration\. Among tree methods, EPIG is more stable and better aligned than EPTree\-style entropy branching\. Flat GRPO aligns best with the GRPO\-style reference, highlighting that topology and advantage construction must be evaluated separately\. Error bars are SEM over three seeds\.
### 7\.4Single\-turn LLM math: tree credit helps, placement ties

On Qwen\-class GSM8K\-hard and MATH\-small, tree\-local training beats flat GRPO, but branch\-placement variants fall within single\-seed noise \(Figure[5](https://arxiv.org/html/2609.20004#S7.F5)\)\. EPIG wins MATH\-small final Pass@1 \(0\.382 versus 0\.206 for flat GRPO\), while EPTree wins GSM8K\-hard \(0\.828\)\. We therefore do not claim that EPIG beats EPTree in single\-turn math; the robust finding is that tree\-local credit is useful and that token masking can dominate branch placement\.

The most consequential LLM engineering lesson was a loss\-mask bug: an incorrect branch\-segment mask trained only a short window and missed answer tokens in thinking\-mode outputs\. Fixing the mask moved EPIG on GSM8K\-hard from 0\.547 to 0\.781 Pass@1 and frozen gradient cosine from 0\.428 to 0\.691, and moved uniform\-tree from 0\.625 to 0\.781, while the two non\-mask methods barely moved \(Figure[6](https://arxiv.org/html/2609.20004#S7.F6)\)\. For LLM tree RL, topology and credit assignment are coupled: a good branch score is useless if its advantage is applied to the wrong tokens\.

![Refer to caption](https://arxiv.org/html/2609.20004v1/figures/fig5.png)Figure 5:Single\-turn LLM math\. Tree methods beat flat GRPO, especially on the harder MATH\-small, but branch placement is within single\-seed noise\. This motivates stateful multi\-turn benchmarks rather than more saturated single\-turn math\.![Refer to caption](https://arxiv.org/html/2609.20004v1/figures/fig6.png)Figure 6:The loss\-mask fix is the large LLM\-side signal\. The two mask\-using methods \(EPIG, uniform tree\) jump on both Pass@1 and frozen gradient cosine, while the non\-mask methods \(EPTree, flat GRPO\) are essentially unchanged on Pass@1, isolating the mask as the cause\.
### 7\.5Multi\-turn Wordle: frozen calibration

Single\-turn math distributes value over many reasoning tokens, making branch placement hard to isolate; multi\-turn games provide clearer states and actions\. On frozen Wordle states from a TextArena\-style environment with a Wordle\-tuned Qwen\-family model, we estimate high\-budget reference values and compare budgeted estimators\. EPIG\-grad achieves the lowest value MSE, 0\.523, versus 0\.568 for entropy and 1\.094 for uniform \(Figure[7](https://arxiv.org/html/2609.20004#S7.F7)\)\. Crucially, pure value variance is worse \(1\.389\): score weighting matters, supporting the theorem’sψ​Q\\psi Qform over a plainVar⁡\[Q\]\\mathrm\{Var\}\[Q\]heuristic\.

![Refer to caption](https://arxiv.org/html/2609.20004v1/figures/fig7.png)Figure 7:Frozen Wordle state calibration\. EPIG\-grad has the lowest value MSE, and the gap to pure value variance shows that gradient/score weighting—not value dispersion alone—drives the advantage\.
### 7\.6Online multi\-turn Wordle

The decisive test is whether the dense\-control and frozen\-calibration advantage survives online, where the policy moves and the action space is large\. We run online Wordle—a Wordle\-tuned 1\.7B model, four methods×\\timestwo seeds, 300 policy updates—training each for the full update budget and logging win rate against environment interactions\.

The Wordle result is the clearest online evidence in the paper \(Figure[8](https://arxiv.org/html/2609.20004#S7.F8)\)\. All three tree methods start below flat GRPO in the first∼30\\sim 30updates—tree advantages are noisier early—but flat GRPO saturates near a win rate of 0\.790 by update 50 and improves no further\. The tree methods keep climbing and overtake it: EPIG\-grad crosses flat GRPO around update 140 and separates over the second half of training, finishing at 0\.850, compared with EPTree \(0\.825\), uniform turn tree \(0\.805\), and flat GRPO \(0\.790\)\. The final EPIG margin over the strongest tree baseline is \+0\.025 win rate, and EPIG gains the most from extended training\.

![Refer to caption](https://arxiv.org/html/2609.20004v1/figures/fig8.png)Figure 8:Online Wordle win rate per policy update \(smoothed, 95% CI over two seeds\)\. The tree methods start below flat GRPO but overtake it as it saturates near 0\.790\. Final win rates are annotated at right\.

## 8Discussion

The experiments support a nuanced conclusion\. In theorem\-aligned settings—especially cloned\-state dense continuous control—EPIG allocation is strongly supported: it reconstructs the reference gradient direction where entropy mis\-points\. In frozen LLM diagnostics it improves tree\-gradient alignment and Wordle value\-MSE\. In online single\-turn LLM math, tree\-local credit is helpful but branch placement is not the bottleneck; token loss masks and advantage construction dominate\. In online multi\-turn Wordle, where the policy moves and the action space is large, EPIG’s allocation advantage reappears as a higher final plateau: it overtakes a fast\-but\-saturating flat GRPO and finishes ahead of every baseline\.

What the theorem does not say\.EPIG is not a universal exploration algorithm\. If the action space is tiny, uniform coverage is already near\-optimal\. If rewards are sparse and pilot rollouts cannot reveal value differences, pure value\-information allocation can be too exploitative\. If branch advantages are applied to irrelevant tokens, topology does not matter\. These are boundary conditions, not contradictions, and our control losses fall exactly where the boundary conditions predict\.

Why this matters for GRPO\.GRPO remains a powerful baseline, but in stochastic, multi\-turn, or recovery\-heavy environments its outcome\-only estimator is inherently coarse\. Trees expose local conditional values, and EPIG provides a principle for deciding which local counterfactuals are worth buying\.

## 9Limitations

Several LLM diagnostics use a single seed and should be read as preliminary rather than as ranked benchmarks; we mark single\-seed comparisons as such throughout and do not overclaim sub\-0\.02 gaps\. The Goodhart and entropy\-drain terms are theoretically motivated and have not been stress\-tested at large proxy\-reward scale\. Finally, EPIG depends on reasonably good pilot estimates of local value and score\-weighted dispersion; sparse\-reward environments may require additional exploration terms\.

## 10Conclusion

We proposed EPIG\-Tree, a compute\-optimal view of rollout\-tree construction\. The core derivation is simple: the variance of the local policy\-gradient contribution splits into decision uncertainty and continuation uncertainty\. New branches should reduce the former; suffix samples should reduce the latter\. This yields explicit allocation laws and shows why entropy is an incomplete branching objective\. The empirical picture matches the theory: EPIG wins cleanly in cloned\-state control and frozen LLM diagnostics, ties cheaper baselines on near\-saturated single\-turn math, and recovers its advantage online in stateful, large\-action multi\-turn games\. A tree is an experimental apparatus for the policy gradient, and EPIG is the design that spends its budget where the gradient is hardest to estimate\.

## Acknowledgments

We thank our colleagues at Tzafon AI for infrastructure support and helpful discussions\.

## References

- \[1\]Z\. Shao et al\., “DeepSeekMath: Pushing the limits of mathematical reasoning in open language models,” arXiv:2402\.03300, 2024\.
- \[2\]Z\. Hou, Z\. Hu, Y\. Li, R\. Lu, J\. Tang, and Y\. Dong, “TreeRL: LLM reinforcement learning with on\-policy tree search,” arXiv:2506\.11902, 2025\.
- \[3\]Y\. Li et al\., “TreePO: Bridging the gap of policy optimization and efficacy and inference efficiency with heuristic tree\-based modeling,” arXiv:2508\.17445, 2025\.
- \[4\]R\. S\. Sutton, D\. McAllester, S\. Singh, and Y\. Mansour, “Policy gradient methods for reinforcement learning with function approximation,” in*Advances in Neural Information Processing Systems*, 1999\.
- \[5\]J\. Schulman, F\. Wolski, P\. Dhariwal, A\. Radford, and O\. Klimov, “Proximal policy optimization algorithms,” arXiv:1707\.06347, 2017\.
- \[6\]G\. Cui et al\., “The entropy mechanism of reinforcement learning for reasoning language models,” arXiv:2505\.22617, 2025\.
- \[7\]S\. Wang et al\., “Beyond the 80/20 rule: High\-entropy minority tokens drive effective reinforcement learning for LLM reasoning,” arXiv:2506\.01939, 2025\.
- \[8\]L\. Gao, J\. Schulman, and J\. Hilton, “Scaling laws for reward model overoptimization,” in*Proceedings of the 40th International Conference on Machine Learning*, PMLR 202, 2023, pp\. 10835–10866\.
- \[9\]F\. Bickford Smith, A\. Kirsch, S\. Farquhar, Y\. Gal, A\. Foster, and T\. Rainforth, “Prediction\-oriented Bayesian active learning,” in*Proceedings of the 26th International Conference on Artificial Intelligence and Statistics*, Proceedings of Machine Learning Research, vol\. 206, 2023, pp\. 7331–7348\.[https://proceedings\.mlr\.press/v206/bickfordsmith23a\.html](https://proceedings.mlr.press/v206/bickfordsmith23a.html)

## Detailed derivations

### Group\-relative baselines and why local values help

Ignoring standard\-deviation normalization, the GRPO group baseline for sampleiiisRi−R¯R\_\{i\}\-\\bar\{R\}\. IfR¯\\bar\{R\}includesRiR\_\{i\}, the estimator is shrunk by a factor close to\(1−1/K\)\(1\-1/K\)in the simplest independent setting; a leave\-one\-out baseline removes this shrinkage\. The larger issue is variance and credit assignment: the same scalar multiplies every score term on the trajectory\. Tree\-local values approximateQ⁡\(h,a\)−V⁡\(h\)Q\(h,a\)\-V\(h\)at internal edges, reducing the chance that recovery actions inherit the sign of unrelated terminal events\.

### Connection to mutual information

The action channelA→YA\\to Yhas utility information

I⁡\(A;Y∣h\)=H⁡\(Y∣h\)−𝔼a∼πh​\[H⁡\(Y∣h,a\)\]\.I\(A;Y\\mid h\)=H\(Y\\mid h\)\-\\mathbb\{E\}\_\{a\\sim\\pi\_\{h\}\}\[H\(Y\\mid h,a\)\]\.IfY\|aY\\mid ais approximately Gaussian with common varianceσ2\\sigma^\{2\}and nearby meansqaq\_\{a\}, then

I⁡\(A;Y∣h\)≈12​σ2​Vara∼πh​\[qa\]\.I\(A;Y\\mid h\)\\approx\\frac\{1\}\{2\\sigma^\{2\}\}\\mathrm\{Var\}\_\{a\\sim\\pi\_\{h\}\}\[q\_\{a\}\]\.The gradient version replacesqaq\_\{a\}withψh​\(a\)​qa\\psi\_\{h\}\(a\)q\_\{a\}\. Thus value dispersion is a small\-noise approximation to reward information, while EPIG\-grad is a small\-noise approximation to gradient information\.

### Occupancy correction

Tree expansion can overrepresent rare prefixes\. If a tree allocates many leaves to a prefix reached by only one root chain, a naive tree advantage gives that prefix too much loss mass\. Occupancy weighting byμh\\mu\_\{h\}corrects this\. In LLMs, exact occupancy is hard to estimate because token prefixes are rarely revisited exactly; practical estimators include root\-chain visitation counts, depth discounting, or segment\-level grouping\.

## Experiment details

The control sweep spans thirteen Gym/MuJoCo environments at branch budget 8 with a 32\-suffix\-per\-action reference; EPIG/value\-information wins gradient\-MSE in 9/13 \(all nine dense continuous\-control\), recovering the reference direction at cosine 0\.998–1\.000\. The frozen GSM8K calibration uses Qwen3\-8B over 64 prompts and 3 seeds against a 64\-leaf reference\. The single\-turn math runs use Qwen\-class models on GSM8K\-hard and MATH\-small\. The classic\-RL ablation comprises a 47\-environment flat\-GRPO baseline and a seven\-environment stochastic\-branchingpp\-ablation \(21 values×\\times3 seeds\)\. The loss\-mask fix described above and ak1→k3k\_\{1\}\\to k\_\{3\}KL\-estimator change were applied before the online multi\-turn runs\.

## Recommended replication protocol for the online test

For independent replication of the online multi\-turn result we recommend: methods flat GRPO, uniform turn tree, EPTree entropy, and EPIG\-grad; a Wordle\-tuned model with headroom; 16 leaves per prompt,K0=4K\_\{0\}=4root chains, and a fixed token/env\-step budget; at least three seeds over the full update budget; and metrics covering reward AUC versus environment interactions, win rate, invalid action rate, KL, entropy, gradient norm, and frozen calibration at checkpoints\.