Research map / Reinforcement learning / Policy gradient and actor-critic

Policy Gradient with Tree Expansion

ICML 2025 · Reinforcement learning › Policy gradient and actor-critic · approach: Variance-reduced policy gradients · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Introduces SoftTreeMax, a softmax generalization whose logits are multi-step tree-expanded trajectory returns, proves gradient variance decays with tree depth at rate set by the second eigenvalue of the behavior transition matrix, and shows a GPU-parallel implementation reduces PG variance by ~1000x on Atari versus distributed PPO.

Problem

Policy gradient methods suffer from high gradient variance and poor sample complexity, and it was unclear how planning via tree expansion could be integrated into the policy parameterization while controlling this variance.

Approach

Replace softmax logits with the expected discounted cumulative (or exponentiated) return over a depth-d tree expansion under a behavior policy, with a learned value at the leaves; the gradient variance then decays as |lambda_2|^d of the behavior chain, and approximate forward models add bias proportional to model error.

Key results

• SoftTreeMax achieves significantly higher reward than distributed PPO across all tested Atari games, obtaining up to 5x more reward.
• Gradient variance of SoftTreeMax is reduced by about three orders of magnitude relative to PPO across games.
• Theoretical bound shows C-SoftTreeMax gradient variance decays as |lambda_2(P^pi_b)|^(2(d-1)), and approximate-model gradient bias is O(d gamma^-d S beta^2 eps).

Measured effects reported in the paper

Surprising findings

• E-SoftTreeMax performed almost identically to C-SoftTreeMax on Atari because the quasi-deterministic dynamics make trajectory logits nearly invariant.
• Deterministic behavior dynamics (e.g., permutation transition matrix) give no exponential variance decay, so the variance benefit can fail.

Limitations

• Analysis restricted to finite state and action spaces; continuous extension only discussed.
• Relies on GPU-based exact simulator (Atari-CuLE), limiting applicability to stochastic environments where tree width grows exponentially.
• Exact expectation computed without sampling, so finite-sample variance is not analyzed; E-SoftTreeMax not empirically validated.

Experimental setup

Datasets: Atari games (via Atari-CuLE GPU simulator)

Benchmarks: Atari games

Metrics: episode reward gradient variance

Baselines: PPO

Builds on: PPO Atari-CuLE softmax policy gradient UCT

Closest papers by meaning

• On the Global Convergence Rates of Softmax Policy Gradient Methods (ICML 2020 · similarity 0.88)
• Beyond Softmax and Entropy: Convergence Rates of Policy Gradients with $\boldsymbol{f}$-SoftArgmax Parameterization $\&$ Coupled Regularization (ICLR 2026 · similarity 0.88)
• Accelerated Policy Gradient: On the Convergence Rates of the Nesterov Momentum for Reinforcement Learning (ICML 2024 · similarity 0.86)
• SPOT: Scalable Policy Optimization with Trees for Markov Decision Processes (NeurIPS 2025 · similarity 0.86)
• SAPG: Split and Aggregate Policy Gradients (ICML 2024 · similarity 0.86)
• Convergence Analysis of Policy Gradient Methods with Dynamic Stochasticity (ICML 2025 · similarity 0.85)