SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Yangqing Fu
Ming Sun
Buqing Nie
Yue Gao*
MoE Key Lab of Artificial Intelligence, AI Institute, Shanghai Jiao Tong University {frank79110, mingsun, Niebuqing, yuegao}@sjtu.edu.cn
Monte Carlo Tree Search (MCTS) algorithms such as AlphaGo and MuZero have achieved superhuman performance in many challenging tasks. However, the computational complexity of MCTS-based algorithms is influenced by the size of the search space. To address this issue, we propose a novel probability tree state abstraction (PTSA) algorithm to improve the search efficiency of MCTS. A general tree state abstraction with path transitivity is defined. In addition, the probability tree state abstraction is proposed for fewer mistakes during the aggregation step. Furthermore, the theoretical guarantees of the transitivity and aggregation error bound are justified. To evaluate the effectiveness of the PTSA algorithm, we integrate it with state-of-the-art MCTS-based algorithms, such as Sampled MuZero and Gumbel MuZero. Experimental results on different tasks demonstrate that our method can accelerate the training process of state-of-the-art algorithms with  $10\% - 45\%$  search space reduction.
With the advancement of reinforcement learning (RL), AlphaGo is the first algorithm that can defeat human professional players in the game of Go [1]. With no supervised learning of expert moves in AlphaGo Zero [2] to restructured self-play in AlphaZero [3], until the recent MuZero that conducts Monte Carlo Tree Search (MCTS) in hidden state space with a learned dynamics model [4]. MuZero presents a powerful generalization framework that allows the algorithm to learn without a perfect simulator in complex tasks. Recently, EfficientZero [5] has made significant progress in the sample efficiency of MCTS-based algorithms. This development has opened up new possibilities for real-world applications, including robotics and self-driving.
When dealing with complex decision-making problems, increasing the search depth is necessary to achieve more accurate exploration in the decision space, but this also leads to higher time and space complexity [6, 7]. For instance, MuZero trained for 12 hours with 1000 TPUs to learn the Go game, and for a single Atari game, it needs 40 TPUs for 12 hours of training [4]. One approach to reduce the computation is the state abstraction method, which aggregates states based on a certain similarity measure to obtain a near-optimal policy [8, 9, 10]. State abstraction is a crucial technique in reinforcement learning (RL) that enables efficient planning, exploration, and generalization [11].
To reduce the search space of MCTS, previous studies have shown the potential of specific state abstraction techniques [12, 13, 14, 15]. However, finding the minimum abstract state space in these studies is an NP-Hard problem [16]. To our knowledge, this work is the first to define path transitivity in the formulation of tree state abstraction, which enables the discovery of the minimum abstract state space in polynomial time. Additionally, recent MCTS-based algorithms utilize neural networks to
*Corresponding author.
Code available at https://github.com/FYQ0919/PTSA-MCTS
estimate the value or reward of states, which may lead to errors in aggregating states with deterministic state abstraction functions. To address this issue, we proposed a probability tree state abstraction function that aggregates states based on the distribution of child node values, which enhances the robustness of aggregation and ensures transitivity.
This paper proposes the probability tree state abstraction (PTSA) algorithm to improve the tree search efficiency of MCTS. The main contributions can be summarized as follows: i) A general tree state abstraction is formulated, and path transitivity is also defined in the formulation. ii) The probability tree state abstraction is proposed for fewer mistakes during the aggregation step. iii) The theoretical guarantees of the transitivity and aggregation error bound are justified. iv) We integrate PTSA with state-of-the-art algorithms and achieve comparable performance with  $10\% - 45\%$  reduction in tree search space.
2.1 MCTS-based Methods
MCTS is a rollout algorithm for solving sequential decision problems [17]. The fundamental idea of MCTS is to search for the most promising actions by randomly sampling the search space, and then expanding the search tree based on those actions [18]. The computational bottlenecks arise from the search loop, especially interacting with the real environment model of each iteration.
Combined with deep neural networks, MCTS-based methods have achieved better performance and efficiency in various complex tasks, such as board games [2], autonomous driving [19], and robot planning [20]. The model-based algorithm MuZero [4] predicts the environmental dynamics model for more efficient simulation. Based on MuZero, EfficientZero [5] is proposed for the training with limited data, which achieves super-human performance on Atari 100K benchmarks. However, both MuZero and EfficientZero require high computational consumption when dealing with complex action spaces. To address arbitrarily complex action spaces, the sample-based policy iteration framework [21] is proposed. Sampled MuZero extends MuZero by sub-sampling a small fraction of possible moves and achieves higher sample efficiency with fewer expanded actions and simulations. The experimental results show that planning over the sampled tree provides a near-optimal approximation [21]. To further reduce the number of simulations, Gumbel MuZero utilizes the Gumbel-Top-k trick to construct efficient planning [22].
2.2 State Abstraction
State abstraction is aimed at reducing the complex state space by aggregating the similar states [11]. The original state space  $S$  can be mapped into a smaller abstract state space  $S_{\phi}$  by state abstraction. By grouping similar states together, state abstraction can help to identify patterns and regularities in the environment, which can inform more effective decision-making [12].
There are two main challenges when applying state abstraction to RL problems. The first challenge is to decrease the value loss between  $S$  and  $S_{\phi}$ . With bounded value loss, the approximate state abstractions allow the agent to learn a near-optimal policy with improved training efficiency [11, 23]. The second challenge is to compute the smallest possible abstract state space, which is proven that the computational complexity is NP-hard [16]. The transitive state abstraction [11] is defined to efficiently compute the smallest possible abstract state space. However, most transitive state abstractions with deterministic predicates have low fault tolerance. To improve the robustness of aggregation, we measure the abstraction probability of the state pairs based on the expected value distributions. In addition, some previous studies have analyzed and discussed some specific state abstractions in tree structure [12, 13, 14, 15], but there is no formal definition and analysis for general tree state abstractions. To our knowledge, PTSA is the first method that defines a general formulation of tree state abstractions for deep MCTS-based methods and analyzes the transitivity and aggregation error under balanced search.
In this section, the prerequisites for our proposed method are introduced. We consider an agent learning in a Markov decision process (MDP) represented as  $\langle S, \mathcal{A}, R, \mathcal{T}, \gamma \rangle$ , where  $S$  denotes the
Table 1: Some previous state abstraction functions [23, 11, 29].
Abstractions	Predicate	Transitive
φa*	a1* = a2* ∧ V*(s1) = V*(s2)	yes
φεa*	a1* = a2* ∧ |V*(s1) - V*(s2)| ≤ ε	no
φQ*	maxa|QM*(s1,a) - QM(s2,a)| = 0	yes
φεQ*	maxa|QM*(s1,a) - QM(s2,a)| ≤ ε	no
φQd*	∀a: [Q*(s1,a)/d] = [Q*(s2,a)/d]	yes
state space,  $\mathcal{A}$  denotes the action space,  $R: S \times \mathcal{A} \mapsto \mathbb{R}$  denotes the reward function,  $\mathcal{T}: S \times \mathcal{A} \mapsto \mathbb{P}(\mathcal{S})$  denotes the transition model, and  $\gamma \in [0,1]$  is the discount factor. The goal is to learn a policy  $\pi: S \mapsto \mathbb{P}(\mathcal{A})$  that maximizes the long-term expected reward in the MDP.
3.1 Monte Carlo Tree Search
MCTS-based algorithms typically involve four stages in the search loop: selection, expansion, simulation, and backpropagation. After  $N$  iterations of the search loop, MCTS generates a policy based on the current states. In the selection stage of each iteration, an action is selected by maximizing over UCB. Notably, AlphaZero [2] and MuZero [4], two successful RL algorithms developed based on a variant of Upper Confidence Bound (UCB) [24] called probabilistic upper confidence tree (PUCT) [25], have achieved remarkable results in board games and Atari games. The formula of PUCT is given by Eq. (1):
$$
a ^ {k} = \operatorname {a r g m a x} _ {a} Q (s, a) + c (s) \cdot P (s, a) \frac {\sqrt {\sum_ {b} N (s , b)}}{1 + N (s , a)}, \tag {1}
$$
where  $k$  is the index of iteration,  $Q(s,a)$  denotes the value of action  $a$  in state  $s$ ,  $c(s)$  is a hyperparameter for balancing the value score with the visiting counts  $N(s,a)$ , and  $P(s,a)$  is the policy prior obtained from neural networks.
3.2 State Abstraction in RL
State abstraction methods aggregate similar environment states to compressed descriptions [26], which simplify the state spaces and significantly reduce the computation time [27]. The state abstraction type is formulated as below [10]:
Definition 3.1. (State Abstraction Type) A state abstraction type is a set of functions  $\phi : S \mapsto S_{\phi}$  related to a fixed predicate on state pairs:  $p_M : S \times S \mapsto \{0,1\}$ . When function  $\phi$  aggregates the state pair  $(s_1, s_2)$  in MDP  $M$ , the predicate  $p_M$  must be true:  $\phi(s_1) = \phi(s_2) \Longrightarrow p_M(s_1, s_2)$ .
The conditions for state abstraction are usually strict, which may cause insufficient compression in state spaces. Recent studies have shown that transitive state abstraction [23] is computationally efficient and can achieve near-optimal decision-making, which is defined as:
Definition 3.2. (Transitive State Abstraction) For a given state abstraction  $\phi$  with predicate  $p_M$ , if  $\forall (s_1, s_2, s_3) \in S$  satisfies  $[p_M(s_1, s_2) \wedge p_M(s_2, s_3)] \Longrightarrow p_M(s_1, s_3)$ , the state abstraction  $\phi$  is a transitive state abstraction.
Some previous state abstraction functions are shown in Table 1, which aim to abstract similar states in general reinforcement learning methods. For instance, abstraction  $\phi_{a^*}$  considers the optimal actions and values of states, which is also widely studied in MCTS-based methods [28].
Transitive state abstraction can reduce the computational cost of finding the smallest abstract state space. As shown in Table 1, most approximate state abstractions are not transitive. However, transitivity is challenging to MCTS-based methods due to the tree structure. In MCTS, a path may contain multiple states, so it is necessary to derive the transitivity of the path from the transitivity of the states. Finding the smallest abstract space in the search tree becomes an NP-hard problem if the state abstraction function lacks transitivity in the path.
4 PTSA-MCTS
In this section, we introduce the proposed probability tree state abstraction (PTSA) algorithm, which improves the tree search efficiency for MCTS. As described in Figure 1, PTSA algorithm improves the search efficiency of MCTS in two aspects: one is reducing the search space by abstracting the original search space, and the other is finding the smallest abstract space efficiently by transitive probability tree state abstraction  $\phi_{Q_{\alpha}^{\psi}}$ . The organization of this section is as follows: Subsection 4.1 gives the formulation of general tree state abstraction. Subsection 4.2 presents our proposed probability tree state abstraction. Subsection 4.3 presents the PTSAZero algorithm, which integrates PTSA with Sampled MuZero, and more information can be found in the Appendix. Subsection 4.4 provides the proofs of transitivity in tree structures and bounded aggregation error under balanced exploration.
Figure 1: The overview structure of PTSA algorithm. The original tree search space in MCTS is mapped into a smaller abstract space efficiently by transitive probability tree state abstraction  $\phi_{Q_{\alpha}^{\psi}}$ .
4.1 Tree State Abstraction Formulation
The formulation of our tree state abstraction is provided with a general abstraction operator for the tree structure, which can be utilized with arbitrary state abstraction types. To facilitate the derivation of theorems and properties, the definition of tree state abstraction is given:
Definition 4.1. (Tree State Abstraction)
For a given tree,  $\mathcal{V}$  and  $\mathcal{B}$  denote the node set and path set respectively. A tree state abstraction is a function  $\phi : \mathcal{V} \mapsto \mathcal{V}_{\phi}$  &  $\mathcal{B} \mapsto \mathcal{B}_{\phi}$  with node predicate  $p_{vM}$  and path predicate  $p_{bM}$  on the sibling path pair:
$$
p _ {v M}: \mathcal {V} \times \mathcal {V} \mapsto \{0, 1 \};
$$
$$
p _ {b M}: \mathcal {B} \times \mathcal {B} \mapsto \{0, 1 \}.
$$
In the tree structure, a path is a sequence of nodes and a node denotes the representation of the corresponding state. The path and node predicates abstract corresponding states of the path and node respectively. For a given path pair  $(b_{1}, b_{2})$  with the same length  $l$ , the predicate on this path pair can be decomposed as:
$$
p _ {b M} \left(b _ {1}, b _ {2}\right) = p _ {v M} \left(v _ {0} ^ {b _ {1}}, v _ {0} ^ {b _ {2}}\right) \wedge \dots \wedge p _ {v M} \left(v _ {l - 1} ^ {b _ {1}}, v _ {l - 1} ^ {b _ {2}}\right), \tag {3}
$$
where  $v_{i-1}^b (i = 1, \dots, l)$  is the  $i$ -th node of path  $b$  with length  $l$ . When function  $\phi$  aggregates the path pair  $(b_1, b_2)$  in MDP  $M$ , the predicate  $p_M$  must be true:
$$
\phi \left(b _ {1}\right) = \phi \left(b _ {2}\right) \Longrightarrow p _ {b M} \left(b _ {1}, b _ {2}\right), \tag {4}
$$
and path  $b_{1}$  and path  $b_{2}$  belong to the same abstract cluster.
The tree state abstraction is applied to two search paths of equal length that start from the same parent node. This ensures that each node along the paths has a corresponding potential abstracted node, while preserving the Markov property after aggregation. For instance, the path pair  $(b_{1}, b_{2})$  from different parent nodes  $v_{a}, v_{b}$  can not be aggregated, which violates the Markov property in the state transitions  $v_{a} \to v_{0}^{b_{1}} \to \ldots \to v_{l-1}^{b_{1}}$  and  $v_{b} \to v_{0}^{b_{2}} \to \ldots \to v_{l-1}^{b_{2}}$ .
Previous studies have proven that tree state abstraction can be an efficient approach to reducing branching factors in MCTS [12, 27, 14]. However, finding the smallest abstract state space is an NP-Hard problem in the previous studies [16]. Following the definition of transitive state abstraction [11], we define the path transitivity in the tree state abstraction formulation:
Definition 4.2. (Path Transitivity) For a given tree state abstraction  $\phi$  with predicate  $p_{bM}$ , if  $\forall (b_1, b_2, b_3) \in \mathcal{B}$  satisfies  $[p_{bM}(b_1, b_2) \wedge p_{bM}(b_2, b_3)] \Rightarrow p_{bM}(b_1, b_3)$ , the state abstraction  $\phi$  has path transitivity.
The definition of path transitivity extends the equivalence of state abstraction in tree state space and general RL state space. Tree State abstraction for MCTS maps the original path space  $\mathcal{B}$  into the abstract path space  $\mathcal{B}_{\phi}$ . For non-transitive tree state abstraction, it is necessary to determine whether all possible path pairs belong to the same abstract cluster, and paths may appear repeatedly in path pairs, which requires a massive computation cost to obtain the smallest  $\mathcal{B}_{\phi}$ .
4.2 Probability Tree State Abstraction
State-of-the-art MCTS-based algorithms utilize neural networks to estimate the value or reward of states. However, hard constraints from previous state abstractions can lead to incorrect aggregation during the early training stage. To reduce the probability of states being mapped to the incorrect abstract space due to bias in network prediction, a novel probability tree state abstraction  $\phi_{Q_{\alpha}^{\psi}}$  is proposed in this work:
Definition 4.3. (Probability Tree State Abstraction  $\phi_{Q_\alpha^\psi}$ ) For a given  $\alpha \in [0,1]$  with node predicate  $p_{vM}$  and path predicate  $p_{bM}$ , the aggregation probability of  $\phi_{Q_\alpha^\psi}$  is defined as:
$$
\mathbb {P} \left\{p _ {b M} \left(b _ {1}, b _ {2}\right) = 1 \right\} \triangleq \mathbb {P} \left\{\phi_ {Q _ {\alpha} ^ {\psi}} \left(b _ {1}\right) = \phi_ {Q _ {\alpha} ^ {\psi}} \left(b _ {2}\right) \right\} = 1 - \prod_ {i} ^ {l} \left(1 - \mathbb {P} \left\{p _ {v M} \left(v _ {i} ^ {b _ {1}}, v _ {i} ^ {b _ {2}}\right) = 1 \right\}\right); \tag {5}
$$
$$
\mathbb {P} \left\{p _ {v M} \left(v _ {i} ^ {b _ {1}}, v _ {i} ^ {b _ {2}}\right) = 1 \right\} \triangleq \alpha \left(1 - D _ {J S} \left(\mathbb {P} \left\{Q ^ {\psi} \left(v _ {i} ^ {b _ {1}}, a\right) \right\} \| \mathbb {P} \left\{Q ^ {\psi} \left(v _ {i} ^ {b _ {2}}, a\right) \right\}\right)\right). \tag {6}
$$
where  $\mathbb{P}\{Q^{\psi}(v,a)\} = \frac{exp(Q^{*}(v,a))}{\sum_{j\in\mathcal{A}}exp(Q^{*}(v,j))}$ , and  $D_{JS}$  is the Jensen-Shannon divergence.
$\phi_{Q_{\alpha}^{\psi}}$  encourages nodes that have the same candidate actions with similar value distribution expectations to be aggregated. Using Jensen-Shannon divergence instead of Kullback-Leibler divergence is more advantageous for numerical stability during computation.
4.3 Integration with Sampled MuZero
Our proposed PTSA algorithm can be integrated with state-of-the-art MCTS-based algorithms. The integration includes two main components: offline learning and online searching. Offline learning involves updating the dynamics, prediction, and value networks by sampling trajectories from a buffer. Online searching involves interacting with the environment to obtain high-quality trajectories, similar to MuZero algorithm. Algorithm1 shows how PTSA can be integrated with Sampled MuZero [21] during the searching stage. Compared with the original Sampled MuZero, lines 4-12 describe how to collect all searched paths and update the corresponding node values during the multiple iterations. Lines 13-19 describe how tree state abstraction reduces the search space efficiently. Based on Theorem 4.4, abstracting the most recently searched path is enough to find the smallest abstract space for MCTS-based methods.  $(\phi (b_{i}) = \phi (b_{s}))$  returns a boolean value, where "true" denotes aggregating  $b_{i}$  and  $b_{s}$ . This boolean value is determined by calculating the probability  $\mathbb{P}(\phi (b_i) = \phi (b_s))$  based on Eq.(5) and Eq.(6).
$S_{L}$  list records the searched paths in the current search tree.  $S_{L}.delete(b)$  and  $S_{L}.add(b)$  refer to removing and recording path  $b$  in  $S_{L}$  respectively. The pruning  $(b_{j})$  action denotes removing unique nodes of path  $b_{j}$  compared to the other abstracted path in the search tree.  $h$  denotes the hidden feature of the original real environment state. Algorithm 1 can be generalized to all tree state abstraction functions by replacing  $\phi_{Q_{\alpha}^{\psi}}$ . The selection, expansion, and backpropagation steps are the same with Sampled MuZero [21]. To maintain the balance of  $\sum_{b}N(s,b)$ , the visit count of the aggregate node needs to be accumulated into the corresponding state pair. Furthermore, aggregated nodes with different sets of legal actions could lead to unnecessary exploration of invalid or irrelevant parts of the search space, slowing down the search and potentially reducing the quality of the results. In our implementation, we ensure that the aggregated node only expands legal actions for the abstracted state, thus avoiding any negative impacts caused by illegal actions. An undoing aggregation operation has been considered in 1. As the value estimation becomes more accurate, some previously aggregated
Algorithm 1 PTSAZero
1:	Input: Root node v0, dynamics network dθ, policy network pθ, value network vθ	11:	Expand child nodes: vexpand(π,h,r)
2:	Initialize searched-path list SL = {{{v0}}}	12:	Backpropagation along path bs
3:	for n = 0,1,2... do	13:	SL.add(bs)
4:	Reset v = v0, searching branch bs = [v]	14:	for bi ∈ SL do
5:	Selection with PUCT	15:	if φQψα(bi) = φQψα(bs) then
6:	Add child v' into bs, v = v'	16:	bj = arg min(b.V)
7:	Update hidden state h and reward r:	17:	b∈(bi,bs)
8:	h, r = dθ(v.parent.h,a)	18:	v0.pruning(bj), SL(delete(bj))
9:	Predict value V and policy π:	19:	end for
10:	V, π = vθ(v.h), pθ(v.h)	20:	end for
nodes in the search will no longer be aggregated in the new search. With each new timestep, the value of the search tree nodes is re-evaluated, leading to changes in the following aggregation results.
In Sampled MuZero algorithm, the computational complexity of simulating from the root node can be expressed as  $\mathcal{O}(N_s\cdot (\mathcal{O}(S) + \mathcal{O}(D) + \mathcal{O}(P) + \mathcal{O}(V)))$ .  $N_{s}$  represents the number of simulations, and  $\mathcal{O}(S)$  denotes the computational complexity of the simulation process, which includes selecting children, expanding, and backpropagating. Additionally,  $\mathcal{O}(D)$ ,  $\mathcal{O}(P)$ , and  $\mathcal{O}(V)$  denote the computational complexities of the dynamics network  $\mathrm{d}\theta$ , policy network  $\mathbf{p}\theta$ , and value network  $\mathbf{v}_{\theta}$ , respectively. The dynamics network predicts the next hidden state  $z$  and corresponding reward  $r$  based on the current hidden state  $h$  and action  $a$ . In Algorithm 1, the time complexity is given by  $\mathcal{O}(N_s\cdot (\log_{|\mathcal{A}|}N_s\cdot c_p + \mathcal{O}(S) + \mathcal{O}(D) + \mathcal{O}(P) + \mathcal{O}(V)))$  under balanced search. The balanced search term  $\log_{|\mathcal{A}|}N_s\cdot c_p$  accounts for the exploration of child nodes, where  $|\mathcal{A}|$  represents the number of possible actions, and  $c_{p}$  is a constant controlling exploration. Since tree state abstraction reduces the branching factor, our algorithm enhances the efficiency of selecting child nodes with a smaller  $N_{s}$ . The specific computational time of different methods can be found in Appendix F.
4.4 Theoretical Justification
Following the formulation of the tree state abstractions, the theoretical analysis is conducted from the following perspectives: i) Transitivity; ii) Aggregation error.
4.4.1 Transitivity
To abstract tree paths efficiently, our next result shows the relationship between path transitivity and node transitivity:
Theorem 4.4. For  $\forall (v_{1},v_{2},v_{3})\in \mathcal{V}$  and  $(b_{1},b_{2},b_{3})\in \mathcal{B}$ :
$$
\begin{array}{l}\left[ \left[ p _ {b M} \left(b _ {1}, b _ {2}\right) \wedge p _ {b M} \left(b _ {2}, b _ {3}\right) \right] \Longrightarrow p _ {b M} \left(b _ {1}, b _ {3}\right) \right] \Longleftrightarrow\\\left. \right.\left[\left[ p _ {b M} \left(b _ {1}, b _ {2}\right) \wedge p _ {b M} \left(b _ {1}, b _ {3}\right)\right]\right.\end{array}
$$
$$
\left[ \left[ p _ {v M} \left(v _ {1}, v _ {2}\right) \wedge p _ {v M} \left(v _ {2}, v _ {3}\right) \right] \Longrightarrow p _ {v M} \left(v _ {1}, v _ {3}\right) \right].
$$
Proof. See Appendix A.
Theorem 4.4 indicates that path transitivity is a sufficient and necessary condition for node transitivity. Compared with the previous transitive state abstractions, the proposed  $\phi_{Q_{\alpha}^{\psi}}$  is also transitive for paths as the following proposition given:
Proposition 4.5. The probability of transitivity for  $\phi_{Q_{\alpha}^{\psi}}$  can be computed as:
$$
\mathbb {P} \left\{\left(p _ {b M} \left(b _ {1}, b _ {2}\right) \wedge p _ {b M} \left(b _ {2}, b _ {3}\right) \Longrightarrow p _ {b M} \left(b _ {1}, b _ {3}\right)\right) \right\} =
$$