McTorch

A Monte Carlo Tree Search framework built on Python and PyTorch, with a C++ backend for performance-critical workloads.

Overview

McTorch provides a clean, environment-agnostic MCTS implementation. The design separates the search engine, environment interface, and execution backend so that each piece can be developed and tested independently.

Two backends are supported:

• TorchBackend — pure Python reference implementation, correct by design.
• CppBackend — C++ accelerated backend, built via Cython bindings.

Repository layout

implement/            source code — package, tests, examples, benchmarks
docs/                 all documentation
management/           project management artifacts
scripts/              build and test helper scripts
assets/               diagrams and images

Installation

Requires Python 3.10+, PyTorch 2.0+, and CMake 3.18+.

The C++ extension links against PyTorch at build time, so PyTorch must be installed before building the package. Use --no-build-isolation so the build step can find it:

# 1. Install PyTorch and build tools
pip install torch numpy wheel pybind11

# 2. Build and install mctorch (core only)
pip install -e . --no-build-isolation

# 3. Optional: include pgx game environments + JAX/mctx reference backend
pip install -e ".[pgx]" --no-build-isolation

# 4. Optional: development tools (pytest, ruff, mypy)
pip install -e ".[dev]" --no-build-isolation

Why --no-build-isolation?
pip normally builds in a sandboxed environment that only has the packages listed under [build-system] in pyproject.toml. PyTorch is too large to list there, so the sandbox would never have it. --no-build-isolation tells pip to build in your current environment instead, where torch is already present.

Quick start

import torch
import torch.nn as nn
from mctorch import alphazero_policy, RootFnOutput, RecurrentFnOutput

# Minimal two-headed network (policy + value)
class Net(nn.Module):
    def forward(self, obs):
        ...  # return policy_logits [B, A], value [B]

net = Net()
board = torch.zeros(16, 9)          # 16 games, 9-cell TicTacToe boards

with torch.no_grad():
    logits, values = net(board)

root = RootFnOutput(prior_logits=logits, value=values, embedding=board)

def recurrent_fn(params, actions, embedding):
    new_board, reward, discount = your_step_fn(embedding, actions)
    with torch.no_grad():
        new_logits, new_values = params(new_board)
    return RecurrentFnOutput(reward=reward, discount=discount,
                             prior_logits=new_logits, value=new_values), new_board

out = alphazero_policy(net, root, recurrent_fn, num_simulations=200)
print(out.action)        # [B] best action per game

See implement/examples/ for complete runnable scripts.

Running tests

bash scripts/run_tests.sh

or directly:

pytest implement/tests/ -v

Running benchmarks

# mctorch C++ arena vs raw baseline
python implement/benchmarks/bench_mctorch_native.py

# mctx JAX reference backend
python implement/benchmarks/bench_mctx.py

# All backends side-by-side
python implement/benchmarks/bench_all_backends.py

# pgx game environments (requires .[pgx] extras)
python implement/benchmarks/bench_pgx.py

Documentation

• API reference
• Architecture
• Examples guide
• Development conventions
• Contributing guide
• Execution plan

License

MIT