Array-Based Monte Carlo Tree Search

TL;DR

This paper introduces an array-based implementation of Monte Carlo Tree Search that enhances computational efficiency and scalability, enabling faster simulations and improved search performance on modern processors.

Contribution

The paper presents a novel array-based approach to Monte Carlo Tree Search that reduces branch prediction overhead and doubles scaling efficiency with search depth.

Findings

Achieves up to 2.8x faster performance in simulations

Reduces branch prediction overhead in MCTS implementations

Improves scalability with search depth on pipelined processors

Abstract

Monte Carlo Tree Search is a popular method for solving decision making problems. Faster implementations allow for more simulations within the same wall clock time, directly improving search performance. To this end, we present an alternative array-based implementation of the classic Upper Confidence bounds applied to Trees algorithm. Our method preserves the logic of the original algorithm, but eliminates the need for branch prediction, enabling faster performance on pipelined processors, and up to a factor of 2.8 times better scaling with search depth in our numerical simulations.

Equations

$${\hat{V}}_{\text{aug}}(H\cup a)=\hat{V}(H\cup a)+c\sqrt{\frac{\log N(H)}{N(H\cup a)}}$$

Full text

1 Introduction

Monte Carlo Tree Search (MCTS), provides efficient, anytime, and asymptotic solutions to sequential decision making problems formulated as Markov Decision Processes (MDPs). It has been applied to a broad range of planning problems Browne et al. (2012); Świechowski et al. (2023) from super human performance in games Silver et al. (2016), to real time planning in dynamical systems Rivière et al. (2024) to even to predicting protein structures Bryant et al. (2022). MCTS works by constructing a tree of simulated possible futures while biasing the tree towards actions leading to high rewards Browne et al. (2012); Świechowski et al. (2023). In doing so, the root node’s value estimate converges to the true optimal value, allowing for computationally tractable and iteratively better approximate solutions, even for problems where closed form solutions do not exists or are infeasible to compute.

The classical Upper Confidence bounds applied to Trees (UCT) result provides value convergence at $\mathcal{O}(\log(N)/N)$ Kocsis and Szepesvári (2006), while more recent results for Fixed-Depth MCTS provides value convergence at a polynomial rate $\mathcal{O}(N^{-1/2})$ Shah et al. (2020). Therefore, increasing the number of simulations performed within a computation budget will result in improved value estimates and performance. However, changes to the underlying algorithm will change these theoretical results.

Instead, we present an alternative array-based implementation designed to optimize the branching performance of MCTS on pipelined processors. A key motivation of our work is the difficulty of deploying MCTS on hardware accelerators such as GPUs due to high memory latency and limited branching performance Pharr and Fernando (2005), with software packages often providing limited branching capability at all Bradbury et al. (2018).

Branches are a type of control flow commonly found in programs where the execution path taken depends (or branches) on the result of a previous instruction. One of the most common occurrences of branching is an "if" statement, which splits execution based on a boolean value. Due to limitations of pipelined architectures, which are nearly ubiquitous on modern processors, branches result in decreased computational performance as visualized in Figure 1 below. To leverage the full performance of a CPU (or GPU), instructions (such as adding two numbers) are queued and processed in order through different sections of the processor simultaneously. When a branch occurs, the processor must guess which execution path will be followed next, and load the appropriate instructions. If the guess is incorrect, these instructions must be flushed, and new instructions loaded at the start of the pipeline, resulting in lost computation cycles Ramamoorthy and Li (1977), a problem further exasperated by memory latency Murphy (2007). As a result, despite using the same logic, different algorithm implementations can have significantly different computational performance when deployed to real systems, leading to the development of techniques to avoid expensive instructions Anderson (2005); Han and Abdelrahman (2011); Hijma et al. (2023).

Branching is fundamental to tree searches, which decide which child node to proceed to (or new child to initialize), at each level of the tree, during every simulation. Our alternative implementation of UCT avoids the need for branch prediction, mitigating the performance impact. Our contributions are summarized as follows:

We present algorithms for predictable child selection in MCTS. 2. 2.

We use these algorithms to create an array-based implementation of MCTS which does not require branch prediction 3. 3.

We validate our method through numerical experiments and show we improve scaling with simulation depth by up to a factor of 2.8 over an equivalent tree-based implementation.

2 Related work

Parallelization

A common method to accelerate tree search is through parallelization to take advantage of multi-threaded computing algorithms and processors. The challenge however, is that tree search methods are fundamentally sequential, as each search uses information gained during the previous backpropagation step. The simplest approach is root parallelization Chaslot et al. (2008), where multiple independent tree searches are run. On completion, the value functions at each root node are averaged to determine the overall optimal value Cazenave and Jouandeau (2007). While conceptually simple, this approach is inefficient, as it results in many repeated computations. Leaf parallelization instead parallelizes roll outs from newly added leaf nodes across multiple threads Chaslot et al. (2008); Cazenave and Jouandeau (2007). This minimizes repeat computation, but can only be employed at the end of the tree. Other approaches perform multiple searches within the same tree, but must carefully manage the communication between threads to avoid data corruption, introducing overhead Chaslot et al. (2008). Still other approaches combine several of these techniques, such as by using CPUs for in-tree parallelization, and GPUs for leaf parallelization Rocki and Suda (2011); Barriga et al. (2014); Cai et al. (2021).

Tree search heuristics

Other approaches seek out heuristics to guide the tree search for more efficient exploration. Learned value functions were used by Alpha Go to approximate the outcome of a game from an intermediate position, giving MCTS earlier signals to bias its exploration towards Silver et al. (2016). Similar approaches have been taken in real-time motion planning, where an expert tree search is used to train a neural network enabling more efficient online planning Riviere et al. (2021). Other methods leverage domain specific knowledge to modify the tree search, by limiting actions, setting sub goals, or imposing constraints Świechowski et al. (2023).

While they achieve substantial empirical performance increases, both parallelization and heuristic methods modify the underlying search algorithm, and therefore loose the theoretical guarantees of traditional MCTS search Kocsis and Szepesvári (2006) or more recent analysis Shah et al. (2020). This inspires our approach of optimizing the implementation of MCTS algorithms instead. As we do not need to modify the search algorithm, we can run the original search algorithm and retain the theoretical guarantees or apply our implementation to these parallelization and heuristic methods as well.

Arrays in tree searches

Arrays have previously been proposed to compress tree searches in the context of string matching for natural language processing Aoe (1989). Branchless search algorithms have also previously been identified as desirable in the context of binary search problems Elmasry and Katajainen (2013), to avoid the penalties of branch misprediction on pipelined architectures Ramamoorthy and Li (1977). However, these methods have yet to be applied to MCTS or decision making problems.

3 Preliminaries

MCTS over a MDP consists of action nodes and state nodes, which together make up a single layer of the tree, as shown in Fig. 2 for a simple system with three actions and three states. Each node stores the total number of visits so far, the value or reward at the node, and the node’s children. To simulate down the tree, an action is selected according to the UCT augmented value function:

[TABLE]

where $H$ represents the history in the tree up until the current state node, and $H\cup a$ represents taking action $a$ at this point in the tree. $\hat{V}(\cdot)$ and $N(\cdot)$ represent the value estimates and visit counts for given histories, respectively. The next state is generated by simulating the system forward using the selected action and compared to all the current child state nodes. If no match is found, a new child is initialized. The search continuous until the simulation depth has been reached, then backpropagates up the tree, adjusting the value and visit counts as it goes. Implementing this algorithm with a tree data structure provides a natural method for branching over possible child nodes. However, this branching is hard to predict.

4 Algorithm Overview

Our objective is to construct a version of MCTS with a predictable execution path by using arrays as the primary data structure. This will allow for better processor performance, by avoiding branch prediction failures in the processor pipeline and by leveraging existing optimizations for array operations. As a motivating example, Algorithm 1 demonstrates using mathematical operations to avoid the need for an "if" statement. This form of micro optimization is valuable on processors which feature expensive branching, have limited access to compiler optimizations, or lack conditional move instructions Angelou et al. (2022).

In MCTS, the fundamental instance of branching is deciding whether to add a child node or not. If the selected action (or state), has already been added to the tree, then the previous node is retrieved and the algorithm traverses downwards. Otherwise, a new node must be initialized before continuing. Unfortunately, this occurrence of branching exhibits different behavior, as opposed to simply resulting in different values like in Algorithm 1.

However, this example still gives a clue to how the branching can be made predictable. If we considered the child nodes to be arranged in a list ordered by time of insertion as diagrammed Fig. 3, then matching an existing node is equivalent to finding the index of the matching child within this list, whereas the index past the end of the list represents a new node to initialize. Following the example in Algorithm 1, the index of the matching child is a, b is the new node, and c is now whether or not a match occurred. Note that while the match index, a, may be ill defined when no match exists, there is no issue because its value will be discarded.

When adding a new node to the tree, we may need to perform some initialization steps, or otherwise execute different behavior than when selecting an existing node. It is still possible to avoid guessing which execution path to follow, but it now comes at a price. Instead of performing one set of behavior conditioned on whether the node is a new or existing child, we perform both behaviors every time, and use techniques like that of Algorithm 1 to nullify the operations that should not be performed.

5 Array-Based MCTS Algorithm Details

Our algorithm, organizes the nodes by layer within the tree search, separating action and state nodes into two arrays. Now, adding a node to the tree means adding a new node to the end of the appropriate array, and parent nodes save the index of each child. This modified structure is visualized in Fig 4 for the first six iterations of a search over the simple three action and three state system we considered earlier. The nodes in each layer are now ordered by time of creation, with the children of different parents being intermixed.

Note that the length of each node array increases by layer. At depth 1, the action nodes are fully expanded, as all three actions which can be taken from the root node have been tried. Conversely, four more state nodes can be added to the first layer to reach a total of nine, three for each of the three actions. These nine state nodes can themselves have three children, and so on. This observation that the branching is bounded, if growing geometrically by layer, motivates our hierarchical structure where the data from each layer is stored in separate arrays, as opposed to storing all state and all action nodes within the same array. The smaller layers near the root node can be stored efficiently, increasing the likelihood they are kept within the processor cache, leading to lower latency and faster execution.

5.1 Action Branching

To store the action node children of a state node, we take advantage of the fact that the set of available actions is constant and numbered throughout the tree search. To do so, we consider a vector of child indexes for each individual state node, resulting in a 2D array for each layer. For a given entry, the value represents the location of the child action node in the next layer, the column indicates the parent state node, while the row index indicates the action taken. All entries are are initialized to a flag value, indicating that action has not yet been tried.

Our action selection method is shown in Algorithm 2 below. Given the current depth and state index, we retrieve the indexes of each child action (the corresponding action node column). If unvisited actions exist (indicated by the flag value), untriedAction is true and we sample from the untried actions. If none exist, we retry the action with the best UCT value.

In the above algorithm, depth and curStateIdx represent the current position of the tree search, actionValues, actionVisits, stateVisits, numActionsAtDepth, and childActionNodes represent global tree data arrays we can access by layer, and we use the helper functions anyUnvistedChild, uctValue, and getRandomUntriedAction to determine if unvisited child actions exist, compute the UCT value function (Eq. (1)) for each action, and return a random untried action (or a default value if none exist), respectively. Note we only need to return the selected nextAction, as all other information has been stored in the global tree data arrays.

5.2 State Branching

We use a similar method for the child state nodes. However, this is complicated by the fact that the states are not already ordered like the actions are. To address this, we make two modifications to the array representing the child state nodes from the action node architecture. First, while each column still corresponds to a given parent node, the rows no longer have a consistent meaning. Instead state nodes are added to each column in order of generation. To track the number of state nodes already added to each column, we add a row to the bottom of the array which counts the number of state nodes already present, initialized to zero. This doubles as the index which a new state should be added to.

The other change we make from the child action node array is to assume a max level of state branching. In simple MDPs, way may know this branching exactly. However, in continuous dynamical systems (particularly with unbounded noise) the branching may be unbounded Ragan et al. (2024). For a given number of simulations and noise level, however, we can conservatively estimate a maximum level of branching and verify it empirically. Further, we can also set different levels of assumed branching by search layer, knowing that actions closer to the root will be taken more often, leading to more branching Other approaches could be taken to bound the amount of state branching, including clipping the noise distribution or varying the amount of noise or discretization in different layers of the tree.

The resulting child state node selection algorithm is given in Algorithm 3 below. Given the depth, action index, and a multi-deminsional state generated by simulation, we begin by retrieving the first $N_{S,l}$ (the assumed state branching at layer $l$) child state indexes from the column corresponding to the current action node. Any unassigned states are initialized to point to the last possible state in the layer, which has a value initialized to NaN. This ensures these entries will not match when we compare each child state against the generated state. We set the matchIdx to the child state matching the most dimensions of the given state, however the matchFlag is only set if all dimensions match. As with the action selection, we use this flag to distinguish between whether to write to matchIdx or a new entry in the child state node array, as well as determine the value of nextStateIdx and if we should update the number of state nodes in the layer and the current column of the child state node array.

In the above algorithm, curStateIdx is the index of the action node we are deteriming the child of, generatedState is our generated state to compare against, stateNodes, numStatesAtDepth, and childStateNodes represent global tree data arrays we can access by layer, $\Sigma$ sums the number of matched states, maxIndex() returns the index of the best matching child state, and size() returns the number of elements in the state. There is no issue with non-unique maxima for stateMatches as this only occurs when no exact matches exist. Again we only need to return the selected nextStateIdx, all other information has been stored in the global tree data arrays.

5.3 Array-Based MCTS Implementation

Combining Algorithms 2 and 3 results in our full array-based MCTS algorithm, visualized below in Fig. 5. Each layer consists of 8 arrays, divided between the action and state nodes, with columns representing a single action or state within a layer. The parent state and action arrays store each node’s parent index. Alternatives to Algorithms 2 and 3 can be constructed by searching over the parent arrays, however these require exhaustive searches that scale poorly with the number of nodes in a layer, and return varying numbers of nodes, which could reintroduce the need for branch prediction.

We present the complete pseudocode of our array-based MCTS algorithm below in Algorithm 4. We first initialize the data arrays for each layer of the search to the default values of zero, unless stated otherwise. We then iterated down the tree, using Algorithms 2 and 3 to select the action and state child nodes before backpropagating up the tree in the standard MCTS fashion. This process is repeated the specific number of simulations, and the action with the best value is then returned.

In Algorithm 4, rootState is the initial state, $N$ the number of simulations to perform, maxDepth indicates how many layers are present in the tree, actionSet is the fixed set of actions for all layers, $N_{S,1-\texttt{maxDepth}}$ represents the state branching limits at each depth, const($m$,$n$,$c$) is a function which creates an $m$ by $n$ matrix with each entry set to $c$, and rewardFunction() gives a reward based on the current state. The variables curMaxNodes and curMaxChildNodes capture the maximum possible branching at the current layer and its children. Knowing that we start with a single root node, there are at most size(actionSet) actions below it, and size(actionSet)$*N_{S,1}$ state nodes below that, and so on, seen visually in Fig. 5. We also know that at most one action and one state node can be added per layer, per simulation, leading to an alternative bound on the number of nodes per layer. Finally, we note that there is no discounting in the backpropagation step in Algorithm 4, as we consider the finite horizon setting, but it can easily be added in.

6 Empirical Results

In this section, we validate our array-based implementation of MCTS by solving a continuous MDP. We compare our performance against the baseline of a traditional tree-based implementation, as well as an alternative array-based method which does not separate nodes by layers. For consistency, each search method was implemented in C++17, with the same compiler optimizations and libraries. All experiments were performed on a AMD Ryzen 7 3700X (8-Core Processor, 3593 Mhz) with 32 GB of memory available.

We consider the challenging problem of navigating a three degree of freedom vehicle adapted from our previous work. A further challenge is that tree search must be able to consider far enough into the future to discover that going around the obstacle is optimal, as opposed to a simple gradient-based solution which will get stuck against the obstacle wall. In Fig. 6(A), we present a 12 time step simulation, demonstrating our array-based method can successfully solve this MDP problem. In Fig. 6(B) we show the growth of the array-based search over a single time step, validating that a branching search is being performed. Not that while some simulations result in points inside the obstacle, these receive a large negative penalty, and are not selected by the algorithm.

Next, we consider the wall-clock time performance of our array-based tree search on this MDP problem. In Fig. 7, we present the results of averaging the simulation times of the first 10 time steps of this experiment over 10 trials. We note that across a wide range of total simulations, our method demonstrates consistently better scaling with increased simulation depth than the tree-based method as well as the array-based method with no layer sorting. Comparing linear fits between the tree-based method and the array-based method, we see a slope of 0.017 seconds per layer of the search when $N=5,000$ and 0.225 seconds per layer when $N=50,000$ for the tree-based method, compared with 0.008 and 0.095 seconds per layer for the array-based method. Taking the ratios of each methods’ scaling with search depth, we see the array-based method outperforms the tree-based method by a factor ranging from 2.2 to 2.8 across the wide range of total simulations.

Interestingly, the array-based method with no sorting by layer has the same or worse scaling as the tree-based implementation. Combined with its dramatically worse wall clock time, this suggest that sorting the nodes by layer is necessary for efficient memory access, and may indicate that the most frequently accessed nodes near the root of the search are in fact remaining within the processor cache.

Another general trend observed was that the tree-based search took longer during the later time steps of the experiment, while the array-based searches did not. This may be due to the tree more efficiently storing broader searches with more branching, which typically occurred at the beginning of the experiment, compared to the array-based implementations, which may be faster when the same values are repeatedly accessed in the CPU cache.

7 Conclusion

We have introduced an array-based implementation of MCTS, and demonstrated that it out scales traditional tree-based implementations in challenging continuous MDP settings, while maintaining theoretical correctness and the associated convergence guarantees. Our results present several opportunities for future work.

First, while recomputation is inevitable if we wish to avoid branch prediction, one way to minimize the redundant effort is to identify which layers will be fully expanded by the end of the tree search. A hybrid approach could first expand these layers completely before computing the rest of the tree. On subsequent simulations, these first layers could be efficiently bypassed, and our predictable array-based search applied only to the layers which will not be fully explored.

Next, while GPUs and other hardware accelerators with limited branching capabilities were a motivating application for the development of our method, the results presented here are all run on a CPU. Deploying our method on hardware accelerators is the most obvious opportunity for future work, and we anticipate significant performance improvements can be achieved. Finally, dedicated profiling could be performed to identify other opportunities for algorithmic optimization, particularly to verify our intuition that keeping the portions of the tree close to the root in the processor cache is essential for rapid searches.

Acknowledgments and Disclosure of Funding

We thank Sina Aghli who shared an algorithm to avoid branching which inspired our predictable branching techniques.