Compute-efficient branching improves reinforcement learning for complex tasks
EPIG-Tree: Compute-Optimal Branching for Gradient-Efficient Reinforcement Learning
Machine LearningArtificial Intelligence
Summary
Reinforcement learning methods often treat a whole sequence of decisions as one reward, which can be inefficient. The authors designed a new method called EPIG-Tree that decides where to explore more carefully during learning, focusing effort where it most reduces uncertainty. This approach makes learning more accurate and effective for tasks involving many steps, such as controlling robots or playing games like Wordle. Their experiments show that EPIG-Tree outperforms previous methods in various environments and improves the quality of learning signals in language models.
What this means in practice
- •For robot control engineers: Improve training efficiency and accuracy in continuous-control robot tasks by strategically allocating computation to critical decision branches during learning.
- •For game ai developers: Enhance reinforcement learning for complex multi-turn games by more precise credit assignment along decision branches to achieve higher performance.
Authors
Nikita Khomich, Leopold Hermansson, Ido Hakimi
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.