How to deal with infinite loops in the MCTS search of AlphaTensor when using a transposition table?

Question

In the published version of the AlphaTensor algorithm, there are two mentions of a transposition table:

In addition, a transposition table is used to recombine different action sequences if they reach the exact same tensor. This can happen particularly often in TensorGame as actions are commutative

and

After simulating $N(s)$ trajectories from state $s$ using MCTS, the normalized visit counts of the actions at the root of the search tree $N(s, a)/N(s)$ form a sample-based improved policy. Differently from AlphaZero and Sampled AlphaZero, we use an adaptive temperature scheme to smooth the normalized visit counts distribution as some states can accumulate an order of magnitude more visits than others because of sub-tree reuse and transposition table.

I don't understand how this use of a transposition table do not lead to infinite loops in the MCTS search, or how to avoid them.

For instance, say that a state $A$ has many children $B_1, \cdots, B_n$. Each of these child is associated to a certain prior given by the Neural Network. Let us assume that the child maximizing the UCPT score is $B_1$. The MCTS search will then go to $B_1$ and expand it if it's a leaf node. Now, say that $A$ is one of the children of $B_1$. It may very well happen at some point that $A$ is $B_1$'s child maximizing the UPCT score. So at some point, the algorithm will:

1. Start from $A$, reckon that $B_1$ maximizes the UCPT and go to $B_1$
2. Reckon that $A$ maximizes the UCPT and go to $A$
3. But at that point, $A$ is not a leaf node, its children are known, and in particular, the one maximizing the UCPT score is $B_1$.

Hence the infinite loop. How to deal with this situation? You want to treat $A$ as a child of $B_1$, so that you don't have to query the neural network for its priors once again, to update the visit count, etc. But at the same time, considering it as a non-leaf node may make this loop appear.

One possible solution would be to create a "copied node", such that all of its characteristics are shared with $A$, but its children are treated as new nodes (though they can also be in the transposition table). That way, each time we get at a leaf node, we will expand the tree instead of circling back. But his seems rather artificial. Is this the best way to deal with this problem?

Answer 1

I would like to acknowledge the excellent work by Nebuly in creating this implementation of AlphaTensor. My answer to your question will be based on his implementation, and I will reference some of his code accordingly. If you are keen on understanding the finer details of AlphaTensor, I highly recommend reading through their code.

Just to clarify, the "infinite loop" you mentioned isn't technically an infinite loop. When executing MCTS, you always need to specify the number of simulations you want the algorithm to perform. Moreover, theoretically, as models are expected to improve with each iteration, such loops should occur less frequently. This is because the model learns from its mistakes, and repeating the same states worsens the returned rewards. Hence, the use of a transposition table in AlphaTensor doesn’t exactly address the repetition of states in a loop but rather serves as a tool to economize the inference cost of the Neural Network.

def monte_carlo_tree_search(
   model: torch.nn.Module,
   state: torch.Tensor,
   n_sim: int,
   t_time,
   n_steps: int,
  game_tree: Dict,
  state_dict: Dict,):
  
  #More Code

for _ in range(n_sim):
    simulate_game(model, state, t_time, n_steps, game_tree, state_dict)
# return next state
possible_states_dict, _, repetitions, N_s_a, q_values, _ = state_dict[
    state_hash
]
possible_states = _recompose_possible_states(possible_states_dict)
next_state_idx = select_future_state(
    possible_states, q_values, N_s_a, repetitions, return_idx=True
)
next_state = possible_states[next_state_idx]
return next_state

Here's an excerpt from the MCTS function which demonstrates the use of transposition tables, represented by two dictionaries: game_tree and state_dict. The process begins with a for loop that runs n_sim times (the limit of simulations per Monte Carlo play), executing the simulate_game() function. Essentially, this function plays the TensorGame and updates the dictionaries.

def simulate_game(
model,
state: torch.Tensor,
t_time: int,
max_steps: int,
game_tree: Dict,
states_dict: Dict,
horizon: int = 5,
):
"""Simulates a game from a given state.

Args:
    model: The model to use for the simulation.
    state (torch.Tensor): The initial state.
    t_time (int): The current time step.
    max_steps (int): The maximum number of steps to simulate.
    game_tree (Dict): The game tree.
    states_dict (Dict): The states dictionary.
    horizon (int): The horizon to use for the simulation.
"""
idx = t_time
max_steps = min(max_steps, t_time + horizon)
state_hash = to_hash(extract_present_state(state))
trajectory = []
# selection
while state_hash in game_tree:
    (
        possible_states_dict,
        old_idx_to_new_idx,
        repetition_map,
        N_s_a,
        q_values,
        actions,
    ) = states_dict[state_hash]
    possible_states = _recompose_possible_states(possible_states_dict)
    state_idx = select_future_state(
        possible_states, q_values, N_s_a, repetition_map, return_idx=True
    )
    trajectory.append((state_hash, state_idx))  # state_hash, action_idx
    future_state = extract_present_state(possible_states[state_idx])
    state = possible_states[state_idx]
    state_hash = to_hash(future_state)
    idx += 1

# expansion
if idx <= max_steps:
    trajectory.append((state_hash, None))
    if not game_is_finished(extract_present_state(state)):
        state = state.to(model.device)
        scalars = get_scalars(state, idx).to(state.device)
        actions, probs, q_values = model(state, scalars)
        (
            possible_states,
            cloned_idx_to_idx,
            repetitions,
            not_dupl_indexes,
        ) = extract_children_states_from_actions(
            state,
            actions,
        )
        not_dupl_actions = actions[:, not_dupl_indexes].to("cpu")
        not_dupl_q_values = torch.zeros(not_dupl_actions.shape[:-1]).to(
            "cpu"
        )
        N_s_a = torch.zeros_like(not_dupl_q_values).to("cpu")
        present_state = extract_present_state(state)
        states_dict[to_hash(present_state)] = (
            _reduce_memory_consumption_before_storing(possible_states),
            cloned_idx_to_idx,
            repetitions,
            N_s_a,
            not_dupl_q_values,
            not_dupl_actions,
        )
        game_tree[to_hash(present_state)] = [
            to_hash(extract_present_state(fut_state))
            for fut_state in possible_states
        ]
        leaf_q_value = q_values
else:
    leaf_q_value = -int(torch.linalg.matrix_rank(state).sum())
# backup
backward_pass(trajectory, states_dict, leaf_q_value=leaf_q_value)

The states_dict will contain values such as q_values, N_s_a (the visit counter of the state), actions, children's states, and a repetition map to address the commutative nature of multiplication, which results in many repeated states.

Conclusion

In conclusion, if you are interested in examining the implementation of AlphaTensor, I would urge you to read through the repository. If an infinite loop occurs, whether it's within the transposition table or not, it will eventually terminate. Based on this implementation, the simulation maintains a trajectory of the state hashes with each iteration. If a repeated state is found, it will append a reference to the state dictionary and append a copy to the state in the game tree. Therefore, your initial idea of creating a copy is valid. I hope this information is helpful to you.