mutrimcts 0.1.0

High-performance C++ MuZero MCTS in learned latent space

MuTriMCTS

---

---

MuTriMCTS is a high-performance Python package providing C++ bindings for MuZero Monte Carlo Tree Search (MCTS) operating in learned latent space.

Features

• 🚀 High Performance: C++ core with Python bindings via pybind11
• 🧠 MuZero Algorithm: MCTS in learned latent space (no game engine required during search)
• 🎯 Clean API: Simple Protocol-based network interface
• 📦 Easy Installation: Available via PyPI
• ✅ Well Tested: Comprehensive test suite
• 🔧 Configurable: Flexible search parameters (simulations, CPUCT, discount, Dirichlet noise)

Installation

From PyPI (when published)

pip install mutrimcts

From Source

git clone https://github.com/lguibr/mutrimcts.git
cd mutrimcts
pip install -e .

Development Setup

# Clone and install with dev dependencies
git clone https://github.com/lguibr/mutrimcts.git
cd mutrimcts
pip install -e ".[dev]"

# Run tests
pytest tests/

# Clean build artifacts (if needed)
# rm -rf build/ src/mutrimcts.egg-info/ dist/ src/mutrimcts/mutrimcts_cpp.*.so

Quick Start

import mutrimcts
import numpy as np

# Implement your MuZero network
class MyMuZeroNetwork(mutrimcts.MuZeroNetworkInterface):
    def initial_inference(self, observation):
        """
        observation → (hidden_state, policy, value)
        """
        hidden_state = self.representation(observation)
        policy, value = self.prediction(hidden_state)
        return hidden_state, policy, value
    
    def recurrent_inference(self, hidden_state, action):
        """
        (hidden_state, action) → (next_hidden_state, reward, policy, value)
        """
        next_hidden, reward = self.dynamics(hidden_state, action)
        policy, value = self.prediction(next_hidden)
        return next_hidden, reward, policy, value

# Configure search
config = mutrimcts.SearchConfiguration(
    max_simulations=50,
    max_depth=10,
    cpuct=1.25,
    dirichlet_alpha=0.3,
    dirichlet_epsilon=0.25,
    discount=0.997  # Important for MuZero!
)

# Run MCTS
network = MyMuZeroNetwork()
observation = get_current_observation()
visit_counts, root_value, mcts_policy = mutrimcts.run_mcts(
    observation, network, config
)

# Use results for training and action selection
# - visit_counts: target for policy loss
# - root_value: used in value bootstrapping  
# - mcts_policy: for action selection (proportional to visits)

API Reference

Network Interface

class MuZeroNetworkInterface(Protocol):
    def initial_inference(self, observation: Any) -> tuple[Any, dict[int, float], float]:
        """Returns: (hidden_state, policy_dict, value)"""
        ...
    
    def recurrent_inference(self, hidden_state: Any, action: int) -> tuple[Any, float, dict[int, float], float]:
        """Returns: (next_hidden_state, reward, policy_dict, value)"""
        ...

Search Configuration

config = SearchConfiguration(
    max_simulations=50,      # Number of MCTS simulations
    max_depth=10,            # Maximum search depth
    cpuct=1.25,              # PUCT exploration constant
    dirichlet_alpha=0.3,     # Dirichlet noise alpha
    dirichlet_epsilon=0.25,  # Dirichlet noise weight
    discount=0.997,          # Discount factor (gamma)
    mcts_batch_size=1        # Batch size for network calls
)

MCTS Function

def run_mcts(
    initial_observation: Any,
    network_interface: MuZeroNetworkInterface,
    config: SearchConfiguration
) -> tuple[dict[int, int], float, dict[int, float]]:
    """
    Returns:
        - visit_counts: dict[int, int] - Visit counts per action
        - root_value: float - Root node value estimate
        - mcts_policy: dict[int, float] - Normalized MCTS policy
    """

Project Structure

mutrimcts/
├── src/mutrimcts/              # Python package source
│   ├── __init__.py             # Package exports
│   ├── config.py               # SearchConfiguration
│   ├── mcts_wrapper.py         # Python entry point
│   └── cpp/                    # C++ source code
│       ├── bindings.cpp        # pybind11 bindings
│       ├── mcts.h/.cpp         # MCTS algorithm
│       ├── python_interface.h  # Network interface
│       ├── config.h            # Config struct
│       └── CMakeLists.txt      # Build configuration
├── tests/                      # Test suite
│   └── test_muzero_mcts.py
├── pyproject.toml              # Package metadata
├── setup.py                    # Build script
└── README.md                   # This file

How It Works

MuTriMCTS implements the MuZero algorithm:

1. Initial Inference: Converts raw observation to latent state
2. Tree Search: MCTS in latent space using learned dynamics
3. Recurrent Inference: Predicts next state, reward, policy, value
4. Backpropagation: Discounted value accumulation
5. Result: Visit counts and improved policy for training

Key Differences from AlphaZero

Feature	AlphaZero	MuZero (MuTriMCTS)
Search Space	Real game states	Learned latent states
Game Engine	Required during search	Only at root
State Representation	Actual game state	Hidden state tensor
Rewards	Only at terminal	Predicted per transition
Network Calls	evaluate_state()	initial_inference() + recurrent_inference()

Development

Building from Source

# Install dependencies
pip install pybind11>=2.10 cmake>=3.14

# Build C++ extension
mkdir build && cd build
cmake ../src/mutrimcts/cpp
cmake --build . --config Release

# Copy to package
cp mutrimcts_cpp.*.so ../src/mutrimcts/

Running Tests

pytest tests/ -v

License

MIT License - see LICENSE file for details

Contributing

Contributions welcome! Please:

1. Fork the repository
2. Create a feature branch
3. Make your changes
4. Add tests
5. Submit a pull request

Citation

If you use MuTriMCTS in your research, please cite:

@software{mutrimcts2025,
  author = {Luis Guilherme P. M.},
  title = {MuTriMCTS: MuZero MCTS in Learned Latent Space},
  year = {2025},
  url = {https://github.com/lguibr/mutrimcts}
}

Links

• Repository: https://github.com/lguibr/mutrimcts
• Issues: https://github.com/lguibr/mutrimcts/issues
• PyPI: https://pypi.org/project/mutrimcts/

Acknowledgments

Based on the MuZero algorithm by DeepMind. Optimized for research and experimentation.