Search and Belief-Space Planning › Learned Search Guidance
World Models HandbookAll chapters Search and Belief-Space Planning › Learned Search GuidanceOne section · Sections in this chapter9 / 13 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Everything so far has treated the search as a pure function of the model. Given enough simulations under the relevant consistency assumptions, MCTS can identify a good action; with a weak rollout policy, finite-budget estimates may remain poor. In the earlier experiment, the guided-rollout planner had substantially lower error than the uniform-rollout planner at every tested budget, and the only difference between them was a hand-written heuristic that followed the BFS shortest path.
That heuristic required us to compute a distance map for the maze, which we could do because we knew its finite transition graph. In a larger or learned environment, comparable guidance might be computed from a model, approximated, or learned from data. This section examines the trained-network case and bridges to the decision-centric research lineages in Part VIII.
Compute is scarce, and a prior is a guess about where search should spend it. A useful prior can improve finite-budget decisions. A misleading positive prior can be corrected when exploration and budget are sufficient; a zero prior can block an action under the displayed PUCT rule unless the implementation adds forced visits, noise, or a floor. Prior quality, confidence, model cost, and exploration therefore have to be evaluated together.
Policy Priors and Leaf Values
Three distinct insertion points illustrate how learned guidance can enter search, and they have different costs and failure modes.
In the tree policy, as an action prior. This is the in PUCT. A policy network suggests which actions are worth trying. A useful prior concentrates the search on a smaller subset of actions, so the tree can spend more of its budget on promising branches. A low-prior action can still be explored when its prior is positive, but a zero prior can suppress its exploration bonus entirely unless the implementation adds a floor or explicit initial visit. The prior guides search; by itself, it does not make search safe. In effect, a concentrated and accurate prior can reduce the tree's effective branching factor, an important source of finite-budget savings because branching compounds with depth.
At the leaves, as a value estimate. Long rollouts can dominate per-simulation cost and produce high-variance returns. AlphaZero replaces default-policy rollouts with a network evaluation at newly expanded leaves. This trades trajectory simulation for a forward pass, but it also makes search sensitive to value error away from the network's training distribution. Belief-space planners can use an analogous value estimate over histories, particles, or learned belief representations when that estimate is available. The estimate is available immediately at every leaf without the cost of a rollout, but systematic value error can bias backed-up statistics unless later expansion or terminal evidence corrects it.
As an amortized policy. Once search has run, its root visit counts define a distribution over actions. A policy network can be trained on that distribution so later searches start from a prior shaped by earlier search. AlphaZero uses this policy-improvement loop while training its value head on game outcomes. In the usual amortized setting, one policy evaluation is cheaper than rerunning the search that supplied its target, but its agreement with search depends on training coverage, on optimization, and on model capacity. Improved agreement can reduce the search budget needed for a target quality, but that is an empirical outcome rather than a guaranteed monotonic trend.
The loop has four parts. (1) Run PUCT with a policy head as the prior and a value head for leaf evaluation. (2) Turn the root visit counts into a target distribution, . (3) Train the policy head against and train the value head against the observed game outcome, with regularization. (4) Repeat. Search supplies a policy-improvement target, while self-play outcomes supervise the value estimate that later searches use. The bookkeeping is what makes the loop work: search must be an accurate supervisor and the policy must be an accurate student, or the two will agree on a shared mistake.
Learned Guidance in Practice
Let us train a policy prior for the maze and see what it does to the search. The setup is deliberately small: a one-hot state encoding, a two-hidden-layer MLP, and a supervised target taken from value iteration. A larger system might train on search targets, demonstrations, or offline trajectories rather than an exact tabular solution. The interface remains a function that returns an action distribution for a given state. In our toy setting the network learns a lookup table with extra steps; with structured inputs and incomplete coverage, the same interface can instead generalize across states, provided the learned representation and training data support it.
import torch
import torch.nn as nn
import torch.nn.functional as F
torch.manual_seed(0)
np.random.seed(0)
def optimal_action(state):
qs = []
for a in range(N_ACTIONS):
s2, r, done = step(state, a)
qs.append(r + (0.0 if done else GAMMA * V_STAR[s2]))
return int(np.argmax(qs))
X_states = torch.eye(N_CELLS, dtype=torch.float32)
y_targets = torch.tensor([optimal_action(s) for s in FREE], dtype=torch.long)
policy_net = nn.Sequential(
nn.Linear(N_CELLS, 96),
nn.ReLU(),
nn.Linear(96, 96),
nn.ReLU(),
nn.Linear(96, N_ACTIONS),
)
optimizer = torch.optim.Adam(policy_net.parameters(), lr=4e-3)
training_losses = []
for epoch in range(800):
logits = policy_net(X_states)
loss = F.cross_entropy(logits, y_targets)
optimizer.zero_grad()
loss.backward()
optimizer.step()
training_losses.append(float(loss.detach()))
with torch.no_grad():
raw_policy_probs = torch.softmax(policy_net(X_states), dim=-1).numpy()
# Mix in a small uniform floor. Cross-entropy on this separable lookup task can
# otherwise drive non-target float32 probabilities to numerical zero, which
# would remove their ordinary PUCT exploration bonus.
PRIOR_FLOOR_MIX = 0.02
policy_probs = (
1.0 - PRIOR_FLOOR_MIX
) * raw_policy_probs + PRIOR_FLOOR_MIX / N_ACTIONS
match_rate = float((policy_probs.argmax(axis=1) == y_targets.numpy()).mean())
# Full-batch optimization is deterministic here, so the raw epoch loss is
# already smooth enough to show directly.
training_loss_curve = np.asarray(training_losses)Initial loss 1.3934 -> final loss 0.000000 Agreement with the optimal action: 100.0% of states
The network fits the target. That is expected: the state space has twenty-two elements, so this is essentially a lookup table with extra steps. What matters for PUCT is the distribution the network produces, because the exploration term uses action probabilities rather than only the top label. Cross-entropy training drives this tiny separable problem toward numerically one-hot outputs, so the code mixes in two percent uniform probability after training. The floor preserves the learned top-action ranking while giving every alternative a nonzero prior factor in the PUCT bonus; forced first visits provide a second safeguard.
Let us look at the training curve and at the prior the network assigns to a few representative states.
At the start state the floored prior puts almost all mass on right, which agrees with the tabular target. At , a cell at the bottom of the dead-end arm, it concentrates on up, the target's unique maximizer there. Each non-target action retains probability from the explicit floor. That floor does not guarantee rapid recovery from a bad prior, but it prevents the learned distribution itself from assigning an action exactly zero PUCT bonus.
Now we install the prior and re-run the convergence experiment.
def learned_prior(state, action):
return float(policy_probs[S_IDX[state], action])
GUIDANCE_ERRORS["Learned prior"] = []
GUIDANCE_VISIT_SHARE["Learned prior"] = []
for budget in BUDGETS:
errs, shares = [], []
for s in SEEDS:
tree = mcts(START, n_sims=budget, priors=learned_prior, seed=s)
a = root_choice(tree)
errs.append(abs(tree.children[a].Q - V_STAR[START]))
shares.append(tree.children[OPTIMAL_FIRST_ACTION].N / budget)
GUIDANCE_ERRORS["Learned prior"].append(float(np.mean(errs)))
GUIDANCE_VISIT_SHARE["Learned prior"].append(float(np.mean(shares)))No guidance 0.981 0.982 0.982 0.985 0.986 0.988 0.990 Guided rollouts 0.174 0.203 0.220 0.239 0.251 0.265 0.274 Learned prior 0.968 0.936 0.595 0.289 0.142 0.070 0.035
The learned prior has the lowest error of the three configurations at the larger tested budgets and reaches an error of at 800 simulations. It does not require the BFS distance map at inference, but its labels were generated from converged value iteration over all twenty-two states, so topology entered through supervision. The one-hot input and complete training set make this a memorization experiment; it demonstrates the PUCT interface and finite-budget allocation effect, not generalization to unseen states.
<matplotlib.legend.Legend at 0x11aacbed0>
The two guidance mechanisms enter at different points. The hand-coded rollout heuristic changes returns beyond the represented tree. The learned prior changes selection before those returns are observed. They can be combined, although this chapter evaluates them separately. AlphaZero-style search combines a policy prior with a learned leaf value instead of a default rollout. A biased leaf value can then bias backed-up search statistics, while continued expansion and terminal outcomes may provide additional correction depending on the domain and budget.
This experiment does not test a misleading prior: the network's top action matches the exact tabular target at every training state. A prior learned from incomplete data can misrank actions, and at finite budget PUCT may then under-sample a low-prior alternative. The outcome depends jointly on the prior probabilities, exploration coefficient, forced-visit or root-noise mechanism, value estimates, and available simulations. To test whether search tolerates a misleading prior, use deliberately corrupted or out-of-distribution priors; the oracle-trained result here cannot establish that tolerance.