# SoftTreeMax: Policy Gradient with Tree Search

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

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 20 of its 35 references that are in the map.
- First posted on arXiv 28 Sep 2022.

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

- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor](https://mlsanity.com/paper/1801.01290/) (2018)
- [Dream to Control: Learning Behaviors by Latent Imagination](https://mlsanity.com/paper/1912.01603/) (2019)
- [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)
- [Fully-Convolutional Siamese Networks for Object Tracking](https://mlsanity.com/paper/1606.09549/) (2016)
- [Brax -- A Differentiable Physics Engine for Large Scale Rigid Body Simulation](https://mlsanity.com/paper/2106.13281/) (2021)
- [On the Global Convergence Rates of Softmax Policy Gradient Methods](https://mlsanity.com/paper/2005.06392/) (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 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)
- [Policy Gradient Methods for Reinforcement Learning with Function Approximation and Action-Dependent Baselines](https://mlsanity.com/paper/1706.06643/) (2017)
- [A Hybrid Stochastic Policy Gradient Algorithm for Reinforcement Learning](https://mlsanity.com/paper/2003.00430/) (2020)
- [How to Combine Tree-Search Methods in Reinforcement Learning](https://mlsanity.com/paper/1809.01843/) (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, 1 of 1)

- [Policy Mirror Descent with Lookahead](https://mlsanity.com/paper/2403.14156/) (2024)

## Abstract (arXiv)

Policy-gradient methods are widely used for learning control policies. They can be easily distributed to multiple workers and reach state-of-the-art results in many domains. Unfortunately, they exhibit large variance and subsequently suffer from high-sample complexity since they aggregate gradients over entire trajectories. At the other extreme, planning methods, like tree search, optimize the policy using single-step transitions that consider future lookahead. These approaches have been mainly considered for value-based algorithms. Planning-based algorithms require a forward model and are computationally intensive at each step, but are more sample efficient. In this work, we introduce SoftTreeMax, the first approach that integrates tree-search into policy gradient. Traditionally, gradients are computed for single state-action pairs. Instead, our tree-based policy structure leverages all gradients at the tree leaves in each environment step. This allows us to reduce the variance of gradients by three orders of magnitude and to benefit from better sample complexity compared with standard policy gradient. On Atari, SoftTreeMax demonstrates up to 5x better performance in faster run-time compared with distributed PPO.