SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\mathbb {P} \left\{\phi_ {Q _ {\alpha} ^ {\psi}} (b _ {1}) = \phi_ {Q _ {\alpha} ^ {\psi}} (b _ {2}) \right\} \mathbb {P} \left\{\phi_ {Q _ {\alpha} ^ {\psi}} (b _ {2}) = \phi_ {Q _ {\alpha} ^ {\psi}} (b _ {3}) \right\} \mathbb {P} \left\{\phi_ {Q _ {\alpha} ^ {\psi}} (b _ {1}) = \phi_ {Q _ {\alpha} ^ {\psi}} (b _ {3}) \right\} + \tag {8}
$$
$$
(1 - \mathbb {P} \{\phi_ {Q _ {\alpha} ^ {\psi}} (b _ {1}) = \phi_ {Q _ {\alpha} ^ {\psi}} (b _ {3}) \}) \mathbb {P} \{\phi_ {Q _ {\alpha} ^ {\psi}} (b _ {2}) = \phi_ {Q _ {\alpha} ^ {\psi}} (b _ {3}) \}) (1 - \mathbb {P} \{\phi_ {Q _ {\alpha} ^ {\psi}} (b _ {1}) = \phi_ {Q _ {\alpha} ^ {\psi}} (b _ {2}) \}).
$$
Proof. See Appendix B.
The computational complexity of computing the smallest possible abstract state space for transitive state abstraction is  $\mathcal{O}\left(|\mathcal{S}|^2\cdot c_p\right)$  [11], where  $c_{p}$  denotes the computational complexity of evaluating for a given state pair.
4.4.2 Aggregation Error Bound
Most approximate state abstractions are non-transitive since their cumulative aggregation value errors are unbounded. We define the aggregation error  $\mathrm{E}^{\phi}$  in path set  $\mathcal{B}$  as:
$$
\mathrm {E} ^ {\phi} = \sum_ {b _ {1}, b _ {2} \in \mathcal {B}} | V ^ {\pi_ {\phi}} (b _ {1}) - V ^ {\pi_ {\phi}} (b _ {2}) |. \tag {9}
$$
For instance, consider a commonly used approximate state abstraction  $\phi_{Q_{\varepsilon}^{*}}$  [23]:
$$
p _ {M} \left(s _ {1}, s _ {2}\right) \triangleq \max  _ {a} \left| Q _ {M} ^ {*} \left(s _ {1}, a\right) - Q _ {M} ^ {*} \left(s _ {2}, a\right) \right| \leq \varepsilon . \tag {10}
$$
with the value loss  $\mathcal{L}_{\phi_{Q_{\varepsilon}^{*}}} = V^{*}(s) - V^{\pi_{\phi}Q_{\varepsilon}^{*}}(s)\leq \frac{2\varepsilon R_{max}}{(1 - \gamma)^{2}}$  . The aggregation error can be bounded as  $\frac{6\varepsilon R_{max}}{(1 - \gamma)^2}$  if two paths of length 3 are abstracted.
In Algorithm1, the aggregation times must be less than the current branching factor<sup>2</sup>. The next result extends to the general tree state abstractions in Algorithm1. The path transitivity implies that the abstracted paths present in the current  $\mathcal{B}_{\phi}$  must belong to different abstract clusters, which can give the following theorem:
Theorem 4.6. Considering a general tree state abstraction  $\phi$  with a transitive predicate  $p(\mathcal{L}_{\phi} \leq \zeta)$ , the aggregation error in Alg. 1 under balanced search is bounded as:
$$
\mathrm {E} ^ {\phi} <   \log_ {| \mathcal {A} |} (N _ {s} + 1) \zeta . \tag {11}
$$
If predicate  $p(\mathcal{L}_{\phi} \leq \zeta)$  is not transitive, the aggregation error is bounded as:
$$
\mathrm {E} ^ {\phi} <   \left(| \mathcal {A} | - 1\right) \log_ {| \mathcal {A} |} \left(N _ {s} + 1\right) \zeta . \tag {12}
$$
Proof. See Appendix C.
Theorem 4.6 provides a theoretical guarantee for finding the smallest  $\mathcal{B}_{\phi}$  within aggregation error bound in Algorithm1. At the same time, this theorem also points out that the transitivity has an important effect on the tree structure: as the size of the action space  $|\mathcal{A}|$  increases, the aggregation error upper bound for transitive tree state abstraction will decrease, while the upper bound for non-transitive state abstraction will increase.
In this section, experiments focus on tree state abstraction for computational efficiency improvement. Firstly,  $PTSA$  is integrated with state-of-the-art MCTS-based RL algorithms Sampled MuZero and Gumbel MuZero and evaluated on various RL tasks with 10 seeds. The aggregation percentages are also evaluated, which reflects the search space reduction in different tasks. In addition, the comparison between probability tree state abstraction and other state abstraction functions is also conducted.
5.1 Results on Performance
To demonstrate the improvement in computational efficiency, the PTSA is integrated with Sampled MuZero in Atari and classic control benchmarks, as well as with Gumbel MuZero in the Gomoku benchmark. The performance is compared against several MCTS-based RL algorithms, including (i) MuZero [4], (ii) Sampled MuZero [21], (iii) EfficientZero [5], (iv) Gumbel MuZero [22].
5.1.1 Atari
Atari games are visually complex environments that pose challenges for MCTS-based algorithms[4, 30]. The results for each benchmark and the normalized result are shown in Figure 2, and the shaded intervals represent the standard deviation of the performance across the different random
The maximum branching factor is less than  $|\mathcal{A}|$  and  $|\hat{\mathcal{A}}|$  in MuZero and Sampled MuZero respectively, where  $\hat{\mathcal{A}}$  is the sampled action space
(a) Alien
(b) Boxing
(c) Freeway
(d) Pong
(e) Tennis
(f) Normalized Score
Figure 2: Experiments results on five Atari games (10 seeds) with a normalized score plot. Sampled MuZero with PTSA (PTSAZero) is compared with three state-of-the-art MCTS-based methods: MuZero [4], Sampled MuZero (SMuZero) [21], and EfficientZero [5]. The tasks include Alien, Boxing, Freeway, Pong, and Tennis.  $N$  denotes the number of simulations. The x-axis is the training time, and the y-axis is the episode return w.r.t training time.
seeds. As Gumbel MuZero does not require large simulations for Atari and control tasks, we only compare its performance in the Gomoku game. When the simulation times  $N = 18$ , MuZero and Sampled MuZero fail to converge within the maximum training time in some tasks. Since the search space of MCTS is mapped into the abstract space with a smaller branching factor, PTSAZero ( $\alpha = 0.7$ ) can converge rapidly with fewer simulations. Although EfficientZero improves the sampling efficiency of MuZero, the increased complexity of the network entails more time to converge with the same computational resources. EfficientZero is tested with different simulation times, and the best case in time efficiency with  $N = 30$  is shown. The normalized score is computed by  $s_{norm} = (s_{agent} - s_{min}) / (s_{max} - s_{min})$ , and the normalized time is computed by  $t_{norm} = (t_{agent} - t_{min}) / (t_{max} - t_{min})$ . The experiment results on Atari benchmarks indicate that Sampled MuZero with PTSA can achieve comparable performance with less training time. For MuZero-based algorithms which require massive computational resources and parallel abilities, PTSA provides a more efficient method with less computational cost.
5.1.2 Control Tasks
In the comparison experiments conducted on Gym benchmarks, including classic control and box2d tasks, certain modifications are made to increase the tree search space in control tasks. Specifically, the action spaces of CartPole-v0 and LunarLander-v2 are discretized into 100 and 36 dimensions, respectively. This discretization of the action space necessitates MCTS-based algorithms to run a larger number of simulations and increases the number of sampled actions in Sampled MuZero.
For both SMuZero and PTSAZero, the number of sampled actions is set to 25 in CartPole-v0 and 12 in LunarLander-v2. Figure 3 demonstrates that PTSAZero exhibits superior computational efficiency compared to other methods. In the LunarLander-v2-36 task, PTSAZero significantly improves the training speed of MuZero by a factor of 3.53. These results highlight the effectiveness of PTSAZero in enhancing the efficiency and performance of MuZero-based algorithms in various control tasks.
5.1.3 Gomoku
Gomoku is a classic board game with multi-step search, where the agent is asked to beat an expert opponent on the  $15 \times 15$  and  $19 \times 19$  boards. Gumbel MuZero replaces heuristic mechanisms in original MCTS algorithms for a smaller number of simulations [22]. In board game tasks, we integrate PTSA algorithm with Gumbel MuZero: tree state abstraction  $\phi_{Q_{\alpha}^{\psi}}$  is utilized to reduce the search space, as described in Algorithm 1. The observation of the Gomoku task is a 2-dimensional matrix describing the distribution of the pieces, and the number of sampled actions for SMuZero,
(a) CartPole
(b)Lunarlander
(c) Gomoku-  $15\times 15$
Figure 3: Experiment results of Gym and Gomoku benchmarks (10 seeds). PTSAGZero denotes the Gumbel MuZero with PTSA algorithm.
(d) Gomoku-  $19\times 19$
(a) Atari
(b) Gomoku & Control
Figure 4: The aggregation percentage on paths during the training process on Atari, Control, and Gomoku tasks varies as the network parameters are updated.
Gumbel MuZero, and PTSAGZero is 30. The training returns of different methods w.r.t. training steps are shown in Figure 3, where a value of 1 represents a win and -1 represents a loss. The results demonstrate that Gumbel MuZero can speed up the training process of Sampled MuZero, and PTSA can provide incremental improvement of Gumbel MuZero. In the Gomoku-  $19 \times 19$  task, despite accelerating the convergence speed of the original algorithms, PTSA does not significantly improve the optimal performance of the algorithms. Experimental results also show that MuZero-based algorithms require a larger number of simulations to learn an effective policy due to the increased size of the action space.
5.2 Search Space Reduction
To analyze the abstracted tree search space, Figure 4 shows the aggregation percentage (the average proportion of aggregated paths) during the training process of PTSAZero on Atari, Control, and Gomoku tasks. Results demonstrate that tree state abstraction reduces the original branching factors by  $10\%$  up to  $45\%$ . As the network parameters are updated during the training process, the aggregation percentage tends to stabilize. As the training progresses towards convergence, the aggregation percentage decreases from  $23.7\%$  to  $8.1\%$  in the Freeway task, while in the Gomoku  $19\times 19$  task, it increases from  $10.1\%$  to  $27.8\%$ . The trend of the aggregation percentage is influenced by the initialization of network and the initial values for the nodes. The converged aggregation percentage may vary depending on specific algorithm parameters and task environments, and it does not imply a certain range of branching factors for all tasks. Each task has unique characteristics and complexities, which can influence the effectiveness of different abstraction functions. It should be emphasized that a larger reduction in the state space does not guarantee improved training efficiency. While a higher aggregation percentage indicates a greater reduction in the search space, the aggregation quality of the abstractions also determines the impact on training efficiency. Improving the accuracy and quality of tree state abstraction during the training process is a future challenge that deserves further research. Furthermore, since the search space of MCTS is reduced by tree state abstraction, the search depth of PTSAZero is deeper than that of SMuZero with the same number of simulations. More results of search depth can be found in Appendix F.
Table 2: Speedup comparison of PTSAZero with different state abstraction functions, where Abs denotes the different tree state abstraction functions (the notations are shown in Table 1). All state abstraction functions are evaluated under the same number of simulations and sampled actions.
Abs	Pong	Boxing	Freeway	Tennis	CartPole	Lunarlander	Acrobot	Average
φa*	2.35±0.44	1.8±0.34	2.44±0.46	1.59±0.3	1.96±0.37	3.03±0.57	2.11±0.4	2.18±0.44
φe*	2.75±0.52	1.77±0.33	2.31±0.43	1.44±0.27	2.0±0.38	2.93±0.55	1.61±0.30	2.12±0.53
φQ*	2.29±0.43	2.11±0.4	2.47±0.46	1.82±0.34	1.64±0.31	1.83±0.34	1.95±0.37	2.00±0.26
φQe*	3.01±0.59	1.95±0.39	2.52±0.49	2.22±0.44	1.74±0.35	3.13±0.61	2.01±0.40	2.37±0.50
φQd*	2.81±0.55	1.84±0.37	2.45±0.48	2.02±0.4	2.44±0.48	3.19±0.62	1.91±0.38	2.38±0.46
φQψ	3.25±0.19	2.00±0.21	2.81±0.33	1.92±0.17	2.14±0.10	3.53±0.31	2.29±0.35	2.56±0.39
5.3 Comparison with Other State Abstraction Functions
We compare the effectiveness of our proposed probability tree state abstraction with other state abstraction functions on Atari and gym benchmarks by integrating them into PTSA algorithm. Referring to the speedup evaluation method in Gumbel MuZero [22], Table 2 shows the speedup of PTSAZero with different state abstraction functions compared to the Sampled MuZero. The specific usage and properties of other state abstraction functions can be found in previous works [23, 11, 29]. We adjusted hyperparameters for different state abstraction functions and selected the best values ( $\epsilon$  and  $d$  are set to 0.5 and 0.2, respectively). The results demonstrate that the proposed probability tree state ab
striction  $\phi_{Q_{\alpha}}\psi$  can achieve better speedup performance on average compared to other state abstraction functions. To evaluate the robustness of  $\phi_{Q_{\alpha}^{\psi}}$ , PTSAZero with different state abstraction functions with  $N = 10$  is compared in the Pong task. As shown in Figure 5, SMuZero and MuZero can not learn an effective policy, and PTSAZero with  $\phi_{Q^{*}}^{\varepsilon}$  and  $\phi_{a^{*}}$  may aggregate incorrect states, which leads to a decrease in performance. We also find that the state abstraction function  $\phi_{a^{*}}$  and  $\phi_{Q^{*}}$  have strict requirements for the abstract conditions, resulting in a small search space reduction (  $1.5\%$  and  $2.8\%$  respectively). Compared with other state abstraction functions,  $\phi_{Q_{\alpha}^{\psi}}$  can abstract the original search space more accurately and robustly, thus accelerating the training process.
Figure 5: Results of PTSAZero with different state abstraction functions in the Pong task when  $N = 10$ . Some state abstraction functions cannot accurately abstract the search space with a small number of simulations.
This paper introduces PTSA algorithm for improving the computational efficiency of MCTS-based algorithms. For efficient abstraction in tree search space, we define path transitivity in the formulation of tree state abstraction. Furthermore, we evaluate that the proposed probability tree state abstraction has a better performance compared with previous state abstraction functions. The experimental results demonstrate that PTSA can be integrated with state-of-the-art algorithms and achieve comparable performance with  $10\% - 45\%$  reduction in tree search space. However, the main limitation of the proposed method is that the parameters of some state abstraction functions need to be manually designed to obtain a more accurate abstract state space. Furthermore, selecting an appropriate state abstraction function based on the characteristics of the state space and transition model is also a potential challenge. Further research will be conducted to address these issues for better performance.
This work was supported by the National Natural Science Foundation of China (Grant No. 92248303 and No. 62373242), the Shanghai Municipal Science and Technology Major Project (Grant No. 2021SHZDZX0102), and the Fundamental Research Funds for the Central Universities.