# Efficient Monte Carlo Tree Search via On-the-Fly State-Conditioned Action Abstraction

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

## Key facts

- Cited by 6 papers in the map (data through 14 Oct 2026).
- Builds on 18 of its 44 references that are in the map.
- First posted on arXiv 2 Jun 2024.

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

- [Grad-CAM: Visual Explanations from Deep Networks via Gradient-based Localization](https://mlsanity.com/paper/1610.02391/) (2016)
- [Mastering Atari, Go, Chess and Shogi by Planning with a Learned Model](https://mlsanity.com/paper/1911.08265/) (2019)
- [The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables](https://mlsanity.com/paper/1611.00712/) (2016)
- [Deep Reinforcement Learning at the Edge of the Statistical Precipice](https://mlsanity.com/paper/2108.13264/) (2021)
- [Learning Invariant Representations for Reinforcement Learning without Reconstruction](https://mlsanity.com/paper/2006.10742/) (2020)
- [Mastering Atari Games with Limited Data](https://mlsanity.com/paper/2111.00210/) (2021)
- [Online and Offline Reinforcement Learning by Planning with a Learned Model](https://mlsanity.com/paper/2104.06294/) (2021)
- [Monte-Carlo Tree Search as Regularized Policy Optimization](https://mlsanity.com/paper/2007.12509/) (2020)
- [Vector Quantized Models for Planning](https://mlsanity.com/paper/2106.04615/) (2021)
- [Context-Specific Independence in Bayesian Networks](https://mlsanity.com/paper/1302.3562/) (2013)
- [Leveraging Factored Action Spaces for Efficient Offline Reinforcement Learning in Healthcare](https://mlsanity.com/paper/2305.01738/) (2023)
- [CAMPs: Learning Context-Specific Abstractions for Efficient Planning in Factored MDPs](https://mlsanity.com/paper/2007.13202/) (2020)
- [Reinforcement Learning in Factored Action Spaces using Tensor Decompositions](https://mlsanity.com/paper/2110.14538/) (2021)
- [Efficient Planning in Combinatorial Action Spaces with Applications to Cooperative Multi-Agent Reinforcement Learning](https://mlsanity.com/paper/2302.04376/) (2023)
- [Adaptive Discretization using Voronoi Trees for Continuous POMDPs](https://mlsanity.com/paper/2302.10439/) (2023)
- [On Discovery of Local Independence over Continuous Variables via Neural Contextual Decomposition](https://mlsanity.com/paper/2405.07220/) (2024)
- [GrASP: Gradient-Based Affordance Selection for Planning](https://mlsanity.com/paper/2202.04772/) (2022)
- [Elastic Monte Carlo Tree Search with State Abstraction for Strategy Game Playing](https://mlsanity.com/paper/2205.15126/) (2022)

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

- [STORM: Search-Guided Generative World Models for Robotic Manipulation](https://mlsanity.com/paper/2512.18477/) (2025)
- [MALinZero: Efficient Low-Dimensional Search for Mastering Complex Multi-Agent Planning](https://mlsanity.com/paper/2511.06142/) (2025)
- [NonZero: Interaction-Guided Exploration for Multi-Agent Monte Carlo Tree Search](https://mlsanity.com/paper/2605.00751/) (2026)
- [AUPO -- Abstracted Until Proven Otherwise: A Reward Distribution Based Abstraction Algorithm](https://mlsanity.com/paper/2510.23214/) (2025)
- [Investigating Intra-Abstraction Policies For Non-exact Abstraction Algorithms](https://mlsanity.com/paper/2510.24297/) (2025)
- [Discovering State Equivalences in UCT Search Trees By Action Pruning](https://mlsanity.com/paper/2510.26346/) (2025)

## Abstract (arXiv)

Monte Carlo Tree Search (MCTS) has showcased its efficacy across a broad spectrum of decision-making problems. However, its performance often degrades under vast combinatorial action space, especially where an action is composed of multiple sub-actions. In this work, we propose an action abstraction based on the compositional structure between a state and sub-actions for improving the efficiency of MCTS under a factored action space. Our method learns a latent dynamics model with an auxiliary network that captures sub-actions relevant to the transition on the current state, which we call state-conditioned action abstraction. Notably, it infers such compositional relationships from high-dimensional observations without the known environment model. During the tree traversal, our method constructs the state-conditioned action abstraction for each node on-the-fly, reducing the search space by discarding the exploration of redundant sub-actions. Experimental results demonstrate the superior sample efficiency of our method compared to vanilla MuZero, which suffers from expansive action space.