MCTS Plugin for Claude Code

Monte Carlo Tree Search (MCTS-LLM) implementation for intelligent problem solving in research, planning, and coding tasks.

Based on the MCTS-LLM paper which demonstrates that LLMs can serve as both:

• World Model: Providing commonsense understanding for predicting action outcomes
• Heuristic Policy: Guiding the search toward promising solutions

Installation

Prerequisites

1. Python 3.10+ with pip
2. Install MCP SDK:

   pip install mcp pyyaml

Install Plugin

# Install from local directory
claude plugin install --plugin-dir /path/to/plugin-mcts

# Or test locally without installing
claude --plugin-dir /path/to/plugin-mcts

Skills

The plugin provides the following skills:

Main Orchestration

• /mcts:mcts - Execute full MCTS-LLM algorithm on a problem

Individual Phases

• /mcts:mcts-select - Selection phase using UCB1
• /mcts:mcts-expand - Expansion phase with LLM as world model
• /mcts:mcts-simulate - Simulation/rollout phase with LLM as policy
• /mcts:mcts-backpropagate - Backpropagation phase to update statistics

Dataset Management

• /mcts:mcts-dataset - CRUD operations for prompt templates

MCP Tools

The plugin provides an MCP server with the following tools:

Tree Management

Tool	Description
mcts_init_tree	Initialize a new MCTS tree for a problem
mcts_select	Execute selection phase with UCB1
mcts_expand	Add child nodes to expand the tree
mcts_simulate	Record simulation results
mcts_backpropagate	Propagate rewards back to root
mcts_get_best_path	Extract the best solution path
mcts_get_tree_stats	Get tree statistics

Observations & Beliefs

Tool	Description
mcts_add_observation	Add an observation to knowledge base
mcts_get_observations	Retrieve observations
mcts_update_belief	Update or create a belief with probability
mcts_get_beliefs	Get all current beliefs

Dataset Management

Tool	Description
mcts_dataset_list	List all prompt templates
mcts_dataset_get	Get a specific prompt
mcts_dataset_create	Create a new prompt
mcts_dataset_update	Update an existing prompt
mcts_dataset_delete	Delete a prompt
mcts_dataset_export	Export prompts to file (backup)
mcts_dataset_import	Import prompts from file
mcts_dataset_reset	Reset to default prompts (requires confirmation)

Usage Examples

Basic MCTS Problem Solving

/mcts:mcts How should I architect a microservices system for an e-commerce platform?

Claude will:

1. Initialize an MCTS tree with the problem
2. Iteratively explore different architectural approaches
3. Simulate outcomes for each approach
4. Learn which paths are most promising
5. Present the best solution with confidence scores

Research Investigation

/mcts:mcts What are the potential causes and solutions for high memory usage in a Node.js application?

Planning Task

/mcts:mcts Plan the migration of a monolithic application to microservices

Coding Problem

/mcts:mcts Implement a rate limiter with the following requirements...

Dataset Management

List Available Prompts

/mcts:mcts-dataset list

Create a Custom Prompt

/mcts:mcts-dataset create

Then provide:

• Name
• Category (research/planning/coding/general)
• Template with {variables}
• Description
• Variable list

Export for Backup

/mcts:mcts-dataset export

Always export before importing!

Import Prompts

/mcts:mcts-dataset import /path/to/prompts.json

Reset to Defaults

/mcts:mcts-dataset reset

Warning: This deletes all custom prompts. Type "yes" to confirm.

How MCTS-LLM Works

The Algorithm

1. Selection: Starting from the root, traverse down the tree using UCB1 to balance exploration vs exploitation:

   UCB = Q/N + c * sqrt(ln(parent_N) / N)
   

2. Expansion: At a leaf node, use the LLM as a world model to generate possible next actions and predict their outcomes.

3. Simulation: From the expanded node, simulate a rollout to a terminal state using the LLM as a heuristic policy. Evaluate the outcome.

4. Backpropagation: Update visit counts (N) and value estimates (Q) from the simulated node back to the root.

5. Repeat: Continue for a set number of iterations or until convergence.

6. Extract Solution: The path with highest average reward represents the best solution.

LLM as World Model

The LLM predicts:

• What actions are valid from a given state
• What state results from each action
• Prior probabilities for action quality

LLM as Heuristic Policy

The LLM guides:

• Which actions to take during simulation
• How to evaluate terminal states
• When to stop exploring a branch

Data Storage

Data is stored in ~/.mcts-data/ or the path specified by MCTS_DATA_DIR:

• trees.json - MCTS tree structures
• observations.json - Recorded observations
• beliefs.json - Probability estimates
• prompts.json - Prompt templates

License

MIT

Acknowledgments

Based on the MCTS-LLM paper: "Large Language Models as Commonsense Knowledge for Large-Scale Task Planning"

@inproceedings{
  zhao2023large,
  title={Large Language Models as Commonsense Knowledge for Large-Scale Task Planning},
  author={Zirui Zhao and Wee Sun Lee and David Hsu},
  booktitle={Thirty-seventh Conference on Neural Information Processing Systems},
  year={2023},
  url={https://openreview.net/forum?id=Wjp1AYB8lH}
}