ReMCTS: Reflection-Enhanced Monte Carlo Tree Search for Code Generation

ReMCTS: selection, expansion, execution, reflection and backpropagation over a tree of candidate programs.

TL;DR: LLM-written code often looks right but fails on edge cases and hidden tests, and single-trajectory self-repair keeps patching an early wrong idea. ReMCTS searches over a tree of candidate programs: every candidate is executed in a sandbox, failures become structured reflections that the branch keeps (path memory) and that other branches retrieve (task-local global memory), and an availability-aware evidence reward tells failed checks apart from missing evidence, so the search also works when no tests are given. On HumanEval and MBPP-Sanitized, ReMCTS with visible tests beats direct generation in 8 of 10 model–dataset pairs; controlled sampling, repair, tree-search and memory ablations show where the gains come from and where they stop.

Each iteration runs selection → expansion → execution → reflection → backpropagation:

Conf rewards low-severity reflections with repair hints and positive reward changes on the path; Risk is the highest similarity to high-severity failures in the global memory ($C=1.4$, $\alpha=0.4$, $\beta=0.6$).

Availability-aware evidence: without explicit tests the unit-test item is unavailable, not failed.

git clone https://github.com/MUSE-AI-Lab/ReMCTS.git
cd ReMCTS
python -m venv .venv && source .venv/bin/activate      # Windows: .venv\Scripts\activate
pip install -e ".[dev]"
pytest -q

Python 3.10+ is enough; models are called through any OpenAI-compatible Chat Completions API, so no GPU is needed. The HumanEval-X C++ pilot additionally needs g++ with C++17.

Offline smoke test (the mock client emits placeholder code, so this checks the pipeline, not accuracy):

remcts-run-search --config default --task-id smoke \
    --prompt "Write a Python function solution(x) that returns x." --tests "assert solution(1) == 1"

With an LLM API (the packaged configs use Qwen3-8B on DashScope; omit --tests for the no_tests setting):

export REMCTS_API_KEY="sk-..."
remcts-run-search --config remcts_no_tests --task-id mex --entry-point solution \
    --prompt "Write a Python function solution(nums) that returns the smallest non-negative integer not in nums."

from remcts.config import load_config
from remcts.core.search_engine import SearchEngine
from remcts.core.state import ProblemSpec

config = load_config("remcts_known")        # packaged name or path to your YAML
config.llm_model_name = "qwen3-32b"
result = SearchEngine(config).run(
    ProblemSpec(task_id="demo", prompt="Return the input unchanged.", entry_point="solution",
                tests="assert solution(1) == 1")
)
print(result.best_code, result.accepted, result.stop_reason)

Every search writes outputs/runs/<task>-<time>/ with the effective config.json, the tree (nodes.jsonl), the reflections (memory.jsonl) and summary.json (including token usage). Another provider only needs llm.model_name and llm.base_url in a YAML config; see docs/CONFIGURATION.md.

python scripts/prepare_datasets.py

downloads HumanEval, MBPP / MBPP-Sanitized, EvalPlus and HumanEval-X C++ from the official sources and derives the splits under data/ (not tracked by Git). File hashes, sizes and row counts are recorded in data/dataset_manifest.json. The experiment drivers use:

Put provider-labelled keys in a local apikeys.txt (ignored by Git); tasks are spread over the keys and move to the next key after a rate-limit or timeout error:

qwen: sk-...
deepseek: sk-...

Then check the plan and the estimated cost before calling any API:

remcts-experiment main --dry-run

All commands share --models, --datasets, --workers, --limit, --budget-cny (stop when the estimated spend reaches the budget) and --report-dir. A report directory holds one <label>.json per run (per-task rows with the returned code, hidden-test results and token usage), summary.md / summary.json, and the search traces under runs/. Interrupted runs resume from <label>.progress.json when the same --report-dir is passed again; remcts-experiment status <report-dir> shows the progress. Details are in docs/REPRODUCIBILITY.md.

Hidden solve rates (%) reported in the paper (Table 1); numbers in parentheses are differences to Direct no_tests.

known is a test-guided mechanism validation and is not equivalent to standard pass@1. Hosted models change over time, so reruns can deviate; every report records the model alias, date, configuration and token usage.

src/remcts/
├── core/          # search tree, memory-augmented UCB selection, expansion, stopping rules, search loop
├── memory/        # structured reflections, path memory, task-local global memory
├── reward/        # unit / execution / semantic / risk scores and the no_tests acceptance gate
├── evaluation/    # contract checks, probes, metamorphic properties
├── sandbox/       # Python and C++17 executors
├── llm/           # OpenAI-compatible client, credential files, prompts, response parsing
├── experiments/   # paper experiments: datasets, baselines, ablations, runner, reports (remcts-experiment)
├── config/        # YAML configs of ReMCTS and the tree-search baselines
├── tasks/         # HumanEval / MBPP loaders
├── cli/           # remcts-run-search, remcts-evaluate, remcts-inspect-run
└── storage/ · utils/
scripts/prepare_datasets.py   # benchmark download and split derivation
docs/                         # configuration, reproducibility, data policy
tests/

@inproceedings{wang2026remcts,
  title         = {{ReMCTS}: Reflection-Enhanced Monte Carlo Tree Search for Code Generation},
  author        = {Wang, Huifei and Huang, Xinying and Sun, Yiheng and Yuan, Yifan},
  booktitle     = {Proceedings of the 2026 Conference on Empirical Methods in Natural Language Processing},
  year          = {2026},
  eprint        = {2609.34717},
  archivePrefix = {arXiv}
}

MIT. Benchmarks and models keep their own licenses and terms.