Learning to Search with MCTSnets

1. Introduction

• MCTSnet learns trajectory selection, state evaluation, and backup instead of relying on domain-tailored search rules.Its architecture incorporates MCTS stages inside a neural network and trains the evaluation, backup, and simulation components jointly.Source: §1 · p. 1
• MCTS simulates future trajectories, evaluates leaf states, and backs up those evaluations to update internal tree values.Leaf evaluation may use random rollouts or a value network, followed by propagation along the simulated trajectory.Source: §1 · p. 1
• MCTSnet represents each search node with a memory vector, traverses forward using a learned simulation policy, and propagates leaf information backward with a backup network.The root memory vector produces the final value or action prediction.Source: §1 · p. 1
• Gradient-based training exposes the search architecture to richer learned components while preserving MCTS structure such as model-based computation, iterative local processing, and structured memory.The internal action sequence is nondifferentiable, so the paper introduces an approximate credit-assignment scheme that exploits the graph’s anytime property.Source: §1 · pp. 1–2
• The method is evaluated in Sokoban, where the authors compare the learned search algorithm with model-free and model-based baselines.The study uses this classic planning task to justify design choices and assess search performance.Source: §1 · p. 2

2. Related Work

• Earlier methods learned evaluation functions for a fixed search procedure, typically using scalar-valued evaluations and decoupling learning from the final search decision.Examples include beam search, minimax, and reinforcement- or supervised-learning approaches in chess, checkers, and shogi.Source: §2 · p. 2
• Prior meta-reasoning approaches optimized search parameters or myopic node choices, but did not provide the same fine-grained learned control over search decisions.The cited work includes black-box optimization for alpha-beta parameters and bandit-based action selection in UCT.Source: §2 · p. 2
• Earlier neural learning-to-plan systems used unstructured memory and had only been demonstrated in toy domains rather than domains approaching Sokoban’s complexity.The cited architecture makes complex branching very unlikely because its memory is unstructured.Source: §2 · p. 2
• MCTSnets extend simulation-based neural architectures by adding tree-structured node-specific memory and learning simulation and tree-expansion strategies.This distinguishes them from I2A, which aggregates simulations, and from architectures that roll out fixed policies or implicit transition models.Source: §2 · p. 2
• Unlike prior approaches that aggregate simulations or use fixed-tree expansion, MCTSnets learn how concrete environment trajectories are selected and expanded.The comparison concerns both the representation of search information and the policy governing search expansion.Source: §2 · p. 2

3.1. MCTS Algorithm

• MCTS maintains node statistics through repeated simulations, leaf evaluation, backup updates, and an anytime search policy that increasingly focuses on promising branches.The standard procedure selects actions during forward simulation, evaluates the reached leaf, propagates values along the traversed path, and chooses the root action from accumulated statistics.Source: §3.1 · pp. 2–3

3.2. MCTSnet: search algorithm

• MCTSnet replaces hand-designed search statistics, backup updates, and expansion policy with learned components trained from data.Its learned search generalizes MCTS by learning what statistics to track, how to propagate them, and how to select expansions.Source: §3.2 · p. 3
• Each MCTSnet simulation traverses the environment from the root, embeds a reached leaf into vector statistics, backs those statistics up through visited nodes, and reads an action distribution from the root.The embedding, simulation policy, backup, and readout networks respectively initialize node memory, select actions, update parent memory, and produce the final decision.Source: §3.2 · p. 3

3.3. MCTSnet: neural network architecture

• MCTSnet is a tree-structured memory whose node states persist between visits, enabling skip-connected updates and partial replanning from a previously computed subtree.Node statistics remain constant until revisited, and a subtree can be reused after transitioning to a successor state.Source: §3.3 · p. 4

3.4. Design choices

• The backup network uses a gated residual update, while the simulation policy can incorporate child statistics and a learned action prior into its logits.The gate can selectively ignore subtree information; the structured policy combines parent and child statistics with model-free prior information.Source: §3.4 · p. 4

3.5. Training MCTSnet

• MCTSnet is trained supervisedly to predict labeled actions with cross-entropy, while REINFORCE supplies gradients for stochastic simulation choices and entropy regularization discourages premature policy convergence.The differentiable network path and the simulation distribution are optimized together from the final prediction objective.Source: §3.5 · pp. 4–5

3.6. A credit assignment technique for anytime algorithms

• REINFORCE is unbiased but has high variance because a single final decision depends on O(M log M) to O(M^2) internal simulation decisions.This motivates the baselines and discounted estimators introduced for anytime algorithms.Source: §3.6 · p. 5
• The proposed anytime credit-assignment method uses improvements in intermediate losses and discounted future rewards to trade off bias and variance in learning simulation decisions.It replaces a purely terminal signal with baselined loss differences, while the discount controls emphasis on short-term versus later improvements.Source: §3.6 · pp. 5–6

4.1. Main results

• MCTSnet performs favorably against model-free and model-based baselines, indicating that learned vector-valued statistics exploit information gathered during search.The comparisons use matched parameters or computation in the model-free settings, while MCTSnet and MCTS share the environment model at test time.Source: §4.1 · pp. 6–7
• 84 ± 1% of levels solved with 25 simulations, versus ≈30% for MCTS with 25 simulations; MCTS required 20 more times simulations to reach MCTSnet’s performance.A stronger PUCT-like baseline reached ≈65% with 25 simulations.Source: §4.1 · p. 7

4.2. Learned statistics and how to perform backups

• The gated residual backup network is more stable and accurate than a simple MLP, while the nonresidual version becomes numerically unstable for M > 25.The gated network can reduce subtree-update influence and gradually depart from its identity skip connection.Source: §4.2 · p. 7

4.3. Learning the simulation policy

• Learned simulation policies outperform random simulation, and lower discount values γ < 1 generally perform better than γ = 1 during training.The proposed credit-assignment scheme and modulated policy architecture improve search-policy learning, although a distilled policy prior does not match the best learned policy.Source: §4.3 · pp. 7–8

4.4. Scalability with number of simulations

• Larger simulation budgets let MCTSnets extract additional relevant information and generally achieve better results with fewer training steps.The comparison varies M during both training and evaluation, relying on weight sharing to support arbitrary M ≥ 1.Source: §4.4 · p. 8, Figure 5 · p. 8

• Learned search performed better than standard MCTS using more than an order of magnitude more computation.The study focuses on small searches, closer in scale to plans processed by the human brain than massive game searches.Source: §5 · p. 8
• MCTSnet preserves MCTS structure while learning which nodes to expand, what statistics to store, and how to propagate them.These components are optimized with gradient-based learning within a neural version of MCTS.Source: §5 · p. 8
• The approach assumes access to a simulator, but the model could instead be learned separately or end-to-end and could support imperfect models.The paper also suggests extending the supervised setup to reinforcement learning through policy iteration with MCTS.Source: §5 · p. 8
• More radical departures from existing search algorithms may offer greater flexibility but could make optimization harder.Source: §5 · p. 8

• Sokoban states are represented as 10×10×4 feature maps with separate channels for walls, the agent, boxes, and targets.The experiments use 10 × 10 layouts containing four boxes and targets.Source: A. Architectural choices and… · p. 11

• 250000 training trajectories from distinct levels were generated with vanilla MCTS using a deep value network and approximately 2000 simulations per step.The generator used a pre-trained value network, depth-wise transposition tables, and subtree reuse after each real step.Source: §A.1 · p. 11
• Approximately 92% of generated levels were solved, with solved trajectories averaging 60 steps and unsolved ones stopping after 100 steps.The testing set contains 2500 trajectories.Source: §A.1 · p. 11

• The embedding network uses three residual blocks and maps convolutional activations to a 128-dimensional vector.Each block has two 64-channel 3x3 convolutions, followed by a 1x1 convolution to 32 channels and a linear projection.Source: §A.2 · p. 11
• The readout is a one-hidden-layer 128-unit MLP, while the policy prior uses an embedding-like network with two residual blocks.All hidden-layer nonlinearities are ReLUs, and the policy prior uses 32-channel convolutions.Source: §A.2 · p. 11

A.3. Implementation and Training

• A single TensorFlow forward execution generates the complete search, including multiple simulations, while sparse operators access tree-node and simulation statistics.Tree variables determine the computation graph for each execution.Source: §A.3 · p. 11
• Gradients flow through node statistics for a fixed expanded tree, enabling learning of the backup and embedding networks.Source: §A.3 · p. 11
• MCTSnet training uses asynchronous distributed optimization with batch size 1, 32 workers, and SGD at learning rate 5e-4.Source: §A.3 · p. 11

B. Additional figures

• Figure 6 depicts one MCTSnet forward pass, detailing early simulations and the final root readout that produces an action distribution.Source: Figure 6 · p. 12
• Figure 7 shows simulation m + 1 traversing to a leaf, then updating memory vectors bottom-up along the visited path while leaving other nodes unchanged.Source: Figure 7 · p. 12
• Figure 8 represents each Sokoban frame as four 10x10 planes supplied to the agent.Source: Figure 8 · p. 12