EPIG-Tree: Compute-Optimal Branching for Gradient-Efficient Reinforcement Learning
This research is relevant for anyone working on reinforcement learning or AI safety, as it provides a new approach to improving the efficiency of policy-gradient estimation, which is a fundamental challenge in the field.
This paper introduces EPIG-Tree, a method for reinforcement learning that uses tree-based rollouts to improve gradient estimation efficiency.
It addresses the inefficiency of existing methods like Group Relative Policy Optimization (GRPO) by placing branches where they most reduce uncertainty about the policy gradient per unit of compute.
The method is shown to reduce gradient MSE in cloned-state control tasks and improve performance in online single-turn math and multi-turn Wordle games compared to flat GRPO and entropy branching.
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.
This report is categorized under AITOP's "Models" desk, focusing on US frontier AI labs and Silicon Valley technology leaders. Within our cross-regional analytical framework, it directly maps to AI developer toolchains, coding agents, and engineering workflows and US-China foundation model benchmarks and commercialization comparison. Key entities and technical topics under continuous tracking include reinforcement learning, gradient estimation, tree-based rollout, ai research.
Source attribution: curated from arXiv cs.LG (@Nikita Khomich, Leopold Hermansson, Ido Hakimi).
