# Tensor Implementation of Monte-Carlo Tree Search for Model-Based Reinforcement Learning
## 3 Proposed Implementation
### Data Structure and Notation  
In this section we introduce data structure and the established notations. Our data structure consists of the multiple tensors to store different attributes of MCTS nodes. During the design, we were inspired by the structure of the Q table. Q table is a matrix holding q values. Rows represent states and columns possible actions in environment. Although Q table stores all possible states, we need to store significantly smaller number of states. MCTS method is applied on batch which consists of \(C_{T}\) observations (i.e., running environments) with \(C_{S}\) MCTS simulations per observation. Therefore, we need to store \(C_{T}+C_{S}\times C_{T}\) states. First \(C_{T}\) states are roots obtained from observations via representation function. Next \(C_{S}\times C_{T}\) states will be explored and stored during the MCTS simulations (\(C_{T}\) states per simulation). In our implementation we have to leave the zero row empty for implementation reasons (explained later), so the total number of rows in each tensor is \(C_{R}=C_{T}+C_{S}\times C_{T}+1\). Data structure of proposed method (shown in Figure 1) consists of following tensors:  
* Tensor \(S\)--all states are held in tensor \(S\). The number of dimensions of tensor \(S\) is equal to the number of dimensions of the state \(+1\). The index of the first dimension represents unique IDs that are shared across all tensors. States are added into tensor \(S\) during the initialization (roots) and simulations (explored states) in order that they were visited (roots have indices from 1 to \(C_{T}\), explored states from the first simulation have indices between \(C_{T}+1\) and \(2\times C_{T}\) and so on). * Tensor \(Q\)--q values are stored in tensor \(Q\), similar to the Q table. First index (of the row) is the same unique ID of the node as we mentioned in tensor \(S\). The size of tensor is \(C_{R}\times|A|\). At the beginning of MCTS method, all values in tensor Q are initialized to zeros. During the phase of backpropagation, q values of traversed nodes are updated. * Tensor \(R\)--holds the predicted rewards by dynamics function. It works on the same principle as tensor \(Q\) and has the same size \(C_{R}\times|A|\). Rewards are put into tensor during phase of expansion. * Tensor \(P\)--stores the predicted probabilities computed by prediction function. The size is similar like previous two tensors \(C_{R}\times|A|\). * Tensor \(N\)--unlike the previous tensors composed by real numbers, tensor \(N\) consists of integers. It stores the numbers of visits of nodes (executed combinations of states and actions). The size of tensor is again \(C_{R}\times|A|\). * Tensor \(E\)--the last tensor \(E\) (with the size \(C_{R}\times|A|\)) holds IDs of children nodes. For example, if there is an edge from the parent node \(ID=i\) after taking action \(a\) to the children node \(ID=j\), then \(E(i,a)=j\). If there is no edge between two nodes, the value in tensor is zero. For that reason, zeroth row is empty in each tensor. The complete data structure consists of the above mentioned 6 tensors. The first index of each tensor is the ID of the node (index of the row). In the tensor \(S\), index \(i\) returns \(i\)-th state (of the node with \(ID=i\)) whereas in other tensors \(i\)-th vector is returned instead. E.g., in tensor \(Q\), the index \(i\) returns vector of q values \(Q(i)=[Q(i,0),Q(i,1),\cdots,Q(i,|A|-1)]\). On the other hand, index \(i\) in combination with action index \(a\) returns scalar q value \(Q(i,a)\). In our implementation, we often index via vectors (as a part of tensor operations). E.g., combination of node index \(i\) and vector of actions \(\vec{a}=[a_{0},a_{15},a_{17}]\), where \(a_{0},a_{15}\) and \(a_{17}\) represent integer numbers, applied to tensor \(Q\) results in \(Q(i,\vec{a})=[Q(i,a_{0}),Q(i,a_{15}),Q(i,a_{17})]\). Figure 1: Data structure—the data of each node is stored in 6 tensors. IDs represent indices of nodes (they are not stored as tensor values).