# Policy Gradient with Tree Expansion

arXiv: https://arxiv.org/abs/2301.13236 · HTML: https://mlsanity.com/paper/2301.13236/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 26 of its 56 references that are in the map.
- First posted on arXiv 30 Jan 2023.

## Builds on (most cited first, 26 of 26)

- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [High-Dimensional Continuous Control Using Generalized Advantage Estimation](https://mlsanity.com/paper/1506.02438/) (2015)
- [World Models](https://mlsanity.com/paper/1803.10122/) (2018)
- [Mastering Atari, Go, Chess and Shogi by Planning with a Learned Model](https://mlsanity.com/paper/1911.08265/) (2019)
- [Isaac Gym: High Performance GPU-Based Physics Simulation For Robot Learning](https://mlsanity.com/paper/2108.10470/) (2021)
- [Brax -- A Differentiable Physics Engine for Large Scale Rigid Body Simulation](https://mlsanity.com/paper/2106.13281/) (2021)
- [Why gradient clipping accelerates training: A theoretical justification for adaptivity](https://mlsanity.com/paper/1905.11881/) (2019)
- [On the Global Convergence Rates of Softmax Policy Gradient Methods](https://mlsanity.com/paper/2005.06392/) (2020)
- [Reward-Free Exploration for Reinforcement Learning](https://mlsanity.com/paper/2002.02794/) (2020)
- [Stochastic Variance-Reduced Policy Gradient](https://mlsanity.com/paper/1806.05618/) (2018)
- [Learning to Simulate Dynamic Environments with GameGAN](https://mlsanity.com/paper/2005.12126/) (2020)
- [The Optimal Reward Baseline for Gradient-Based Reinforcement Learning](https://mlsanity.com/paper/1301.2315/) (2013)
- [On the Convergence and Sample Efficiency of Variance-Reduced Policy Gradient Method](https://mlsanity.com/paper/2102.08607/) (2021)
- [Variance Reduction for Policy Gradient with Action-Dependent Factorized Baselines](https://mlsanity.com/paper/1803.07246/) (2018)
- [An Improved Analysis of (Variance-Reduced) Policy Gradient and Natural Policy Gradient Methods](https://mlsanity.com/paper/2211.07937/) (2022)
- [An Improved Convergence Analysis of Stochastic Variance-Reduced Policy Gradient](https://mlsanity.com/paper/1905.12615/) (2019)
- [Softmax Policy Gradient Methods Can Take Exponential Time to Converge](https://mlsanity.com/paper/2102.11270/) (2021)
- [Exponential Bellman Equation and Improved Regret Bounds for Risk-Sensitive Reinforcement Learning](https://mlsanity.com/paper/2111.03947/) (2021)
- [Active Exploration in Markov Decision Processes](https://mlsanity.com/paper/1902.11199/) (2019)
- [On the Global Optimum Convergence of Momentum-based Policy Gradient](https://mlsanity.com/paper/2110.10116/) (2021)
- [Policy Gradient Methods for Reinforcement Learning with Function Approximation and Action-Dependent Baselines](https://mlsanity.com/paper/1706.06643/) (2017)
- [Beyond the One Step Greedy Approach in Reinforcement Learning](https://mlsanity.com/paper/1802.03654/) (2018)
- [A Hybrid Stochastic Policy Gradient Algorithm for Reinforcement Learning](https://mlsanity.com/paper/2003.00430/) (2020)
- [When Simple Exploration is Sample Efficient: Identifying Sufficient Conditions for Random Exploration to Yield PAC RL Algorithms](https://mlsanity.com/paper/1805.09045/) (2018)
- [Improve Agents without Retraining: Parallel Tree Search with Off-Policy Correction](https://mlsanity.com/paper/2107.01715/) (2021)
- [Planning and Learning with Adaptive Lookahead](https://mlsanity.com/paper/2201.12403/) (2022)

## Built on by (most cited first, 3 of 3)

- [Acceleration in Policy Optimization](https://mlsanity.com/paper/2306.10587/) (2023)
- [Off-OAB: Off-Policy Policy Gradient Method with Optimal Action-Dependent Baseline](https://mlsanity.com/paper/2405.02572/) (2024)
- [OPTS-TTPO: Enhancing Finite-Sample Policy-Gradient Learning with Tree Search](https://mlsanity.com/paper/2609.40035/) (2026)

## Abstract (arXiv)

Policy gradient methods are notorious for having 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 analyze SoftTreeMax and explain how tree expansion 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 variance 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 implementation in Atari, we show that SoftTreeMax reduces the gradient variance by three orders of magnitude. This leads to better sample complexity and improved performance compared to distributed PPO.