---

name: mctx description: "Use Mctx for JAX-native Monte Carlo tree search, MuZero/Gumbel MuZero policies, stochastic MuZero search, and search-tree inspection." metadata: disco-role: operating disable-model-invocation: true license: Apache 2.0

Mctx

Mctx is a compact JAX library for batched Monte Carlo tree search. Use this skill when the task asks you to plan with a learned model, improve a policy with search, or inspect the resulting search tree.

Start here

• Install the library with pip install mctx.
• If you want a quick health check, run python scripts/check_install.py.
• If you want a safe runnable example, run python scripts/policy_improvement_demo.py --num_runs=1.
• Read references/api-reference.md when you need exact function signatures or tensor-shape contracts.
• Read references/workflows.md when you need an end-to-end recipe.
• Read references/troubleshooting.md when search output, shapes, or optional dependencies are failing.

What this skill covers

1. Batched search on a learned model

Use search(...) when you already have:

• a batched root state,
• a recurrent model for expansion,
• a root action-selection function,
• an interior action-selection function,
• and a simulation budget.

This is the lowest-level entry point and is best when you need to customize the search loop.

2. Classic MuZero-style planning

Use muzero_policy(...) when you want the standard MuZero policy improvement loop with Dirichlet root noise, PUCT-style action selection, and visit-count based action weights.

3. Gumbel MuZero planning

Use gumbel_muzero_policy(...) when you want the Gumbel MuZero variant and its default Q-value completion strategy. This is the best default when the task mentions policy improvement by planning with Gumbel.

4. Stochastic MuZero planning

Use stochastic_muzero_policy(...) when the environment alternates between decision nodes and chance nodes. Provide separate decision and chance recurrent functions; the wrapper handles the internal state structure.

How to choose the right helper

• Use qtransform_by_parent_and_siblings for the default MuZero-style value normalization.
• Use qtransform_completed_by_mix_value for the default Gumbel MuZero path.
• Use qtransform_by_min_max when you already know a bounded Q-value range.
• Prefer the policy helper that matches the search style instead of reaching into _src internals.

Inspecting results

After search, prefer the public summary APIs over raw tree fields unless you are troubleshooting.

• PolicyOutput.action is the chosen action.
• PolicyOutput.action_weights are the training targets for the policy head.
• PolicyOutput.search_tree holds the final Tree.
• Tree.summary() exposes root visit counts, visit probabilities, value, and Q-values.
• Tree.qvalues(node_index) is the best way to inspect a node.

Common input contract

• Treat action masks as 1 = invalid, 0 = valid.
• Keep root and recurrent outputs batched.
• Make prior_logits and invalid_actions have the same trailing action dimension.
• Keep recurrent_fn outputs aligned with the batch size of the root.
• Use max_depth only when you intentionally want to cap traversal depth.

Bundled scripts

• scripts/check_install.py verifies the installed package, imports the public API, and reports the active JAX backend.
• scripts/policy_improvement_demo.py is a tiny bandit-style smoke demo that shows how planning can improve a policy.

Optional dependency note

This skill does not bundle tree visualization. If you need Graphviz-style tree rendering later, treat it as a separate dependency problem instead of assuming it is available in a minimal install.

When to read the deeper references

• Read references/api-reference.md before wiring a custom model or action mask.
• Read references/workflows.md before writing a new MuZero or Gumbel MuZero driver.
• Read references/troubleshooting.md when output is flat, invalid-action handling looks wrong, or JAX is on the wrong backend.
• Read references/repo-provenance.md when you need to check whether this skill is stale against the source checkout that produced it.