ICML 2025

Policy Gradient with Tree Expansion
Gal Dalal * 1 Assaf Hallak * 1 Gugan Thoppe 2 Shie Mannor 1 3 Gal Chechik 1 4
Abstract
Policy gradient methods are notorious for hav-
ing a large variance and high sample complexity.
To mitigate this, we introduce SoftTreeMax—a
generalization of softmax that employs planning.
In SoftTreeMax, we extend the traditional logits
with the multi-step discounted cumulative reward,
topped with the logits of future states. We an-
alyze SoftTreeMax and explain how tree expan-
sion helps to reduce its gradient variance. We
prove that the variance depends on the chosen
tree-expansion policy. Specifically, we show that
the closer the induced transitions are to being
state-independent, the stronger the variance decay.
With approximate forward models, we prove that
the resulting gradient bias diminishes with the
approximation error while retaining the same vari-
ance reduction. Ours is the first result to bound
the gradient bias for an approximate model. In
a practical implementation of SoftTreeMax, we
utilize a parallel GPU-based simulator for fast and
efficient tree expansion. Using this implementa-
tion in Atari, we show that SoftTreeMax reduces
the gradient variance by three orders of magni-
tude. This leads to better sample complexity and
improved performance compared to distributed
PPO.
1. Introduction
Policy Gradient (PG) methods (Sutton et al., 1999) for Re-
inforcement Learning (RL) are often the first choice for
environments that allow numerous interactions at a fast pace
(Schulman et al., 2017). Their success is attributed to sev-
eral factors: they are easy to distribute to multiple workers,
require no assumptions on the underlying value function,
and have both on-policy and off-policy variants.
*Equal contribution
1NVIDIA Research 2Indian Institute of
Science 3Technion University 4Bar-Ilan University. Correspon-
dence to: Gal dalal <gdalal@nvidia.com>, Assaf Hallak <ahal-
lak@nvidia.com>.
Proceedings of the 42 nd International Confer...