Monte Carlo Tree Search with UCT

The problem

MCTS (Monte Carlo tree search) finds a good first move by trying things out: it builds a tree of positions, plays many random games to the end, and puts more effort where the results look promising. You will implement it for a single-player problem described by three functions:

• actions(state): the list of legal actions, empty when state is terminal (the end);
• step(state, action): the next state;
• reward(state): the final reward of a terminal state, between 0 and 1.

The balance: UCT

At every node, the question is which child to explore next: the one with the best results so far (exploitation), or one tried only a few times, which might be better than it looks (exploration)? The UCT (Upper Confidence bounds applied to Trees) score balances them:

where $W$ is the total reward collected through the child, $N$ its number of visits (so $W/N$ is its average reward), $N_{\text{parent}}$ the visits of its parent, $\ln$ the natural logarithm and $c$ a constant (such as 1.4). A child with $N = 0$ has score float("inf"): every child is tried once.

Example. A parent visited 10 times, $c = 1.4$.

• Child A: $W = 6$, $N = 8$: $0.75 + 1.4\sqrt{\ln 10 / 8} = 0.75 + 0.75 = 1.50$.
• Child B: $W = 1$, $N = 2$: $0.5 + 1.4\sqrt{\ln 10 / 2} = 0.5 + 1.50 = 2.00$.

B has the lower average but is selected: it has only been tried twice. Implement this as uct(W, N, N_parent, c).

The four phases

Keep, for every node, its visit count $N$, its total reward $W$, its children (by action) and its untried actions. Repeat n_iter times:

1. Selection: from the root, while the node has no untried actions and is not terminal, move to the child with the highest uct (the first in action order on a tie).
2. Expansion: if the node has untried actions, take the first one, create its child, and move to it.
3. Simulation: from the new node's state, play random actions, acts[rng.integers(len(acts))], until a terminal state, and take its reward.
4. Backpropagation: add 1 to $N$ and the reward to $W$ of every node on the path from the root to the new node, root included.

Return the root action whose child has the most visits (the first in action order on a tie).

Example. The root has two actions, a (leads to the end with reward 0) and b (reward 1).

• Iteration 1: the root has untried actions a and b; expansion takes the first, a. It is terminal: reward 0. Root: $N = 1$, $W = 0$; child a: $N = 1$, $W = 0$.
• Iteration 2: the root still has an untried action: expansion takes b, reward 1. Root: $N = 2$, $W = 1$; child b: $N = 1$, $W = 1$.
• Iteration 3: no untried actions. $\mathrm{UCT}(a) = 0 + 1.4\sqrt{\ln 2 / 1} = 1.17$, $\mathrm{UCT}(b) = 1 + 1.17 = 2.17$: b is selected.
• After a few more iterations b has most of the visits: mcts returns "b".

What to implement

• uct(W, N, N_parent, c): the score above, or float("inf") if $N = 0$.
• mcts(root, actions, step, reward, n_iter, c, rng): the search, returning the best root action.

Requirements

• Untried actions are tried in the order actions(state) returns them.
• Use rng only in the simulation, exactly as described.