Gomoku

中文说明

AlphaZero-style Gomoku prototype written in Rust, with a Python/PyTorch training pipeline scaffold.

The v1 goal is pipeline correctness rather than playing strength: rules, environment state, MCTS search, self-play samples, training, ONNX export, and eventually Rust-side ONNX inference should all be deterministic and testable.

Rules

• Board: 15x15.
• Ruleset: free-style Gomoku.
• Black moves first.
• No forbidden moves.
• Five or more consecutive stones wins.
• A full board with no winner is a draw.

Current Status

Implemented:

• Pure Rust rules engine with fixed-size board storage.
• Legal placement errors via Result, without panics for invalid moves.
• Incremental win detection from the last move across four directions.
• Gym-like GomokuEnv wrapper.
• Current-player observation encoding with shape [3, 15, 15].
• Legal action mask with shape [225].
• AlphaZero-style PUCT MCTS with pluggable policy/value evaluation.
• UniformEvaluator for bootstrap search and tests.
• RandomAgent and PureMctsAgent.
• In-memory self-play sample generation and Rust .npz shard export.
• Python dataset loader for .npz shards.
• Small residual PyTorch policy/value network.
• Basic training and ONNX export scripts.

Not implemented yet:

• Rust ONNX Runtime-backed evaluator.
• Scriptable evaluation runner.
• egui/eframe GUI.
• End-to-end generated data -> training -> ONNX -> Rust inference loop.

Project Layout

src/
  game.rs       Pure Gomoku rules engine.
  env.rs        Gym-like environment and observation encoding.
  mcts.rs       PUCT search and policy/value evaluator trait.
  agent.rs      Random and pure-MCTS agents.
  selfplay.rs   Self-play sample generation and .npz shard export.
  main.rs       Smoke-game executable.

python/train/
  gomoku_train/dataset.py   Rust .npz shard loader.
  gomoku_train/model.py     Small residual policy/value network.
  train.py                  Training entrypoint.
  export_onnx.py            Checkpoint to ONNX export entrypoint.

Rust Commands

Run the test suite:

cargo test

Run formatting check:

cargo fmt --check

Run a smoke game using pure MCTS:

cargo run

Expected output is similar to:

Gomoku smoke game finished: winner=Some(Black), draw=false, moves=49

Generate a self-play .npz shard:

cargo run -- selfplay --output data/selfplay_0001.npz --games 1 --simulations 32 --seed 0

This writes:

• data/selfplay_0001.npz: training arrays.
• data/selfplay_0001.json: metadata for the shard.

Python Training

The Python environment is managed with uv.

Create or update the training environment:

cd python/train
uv sync

Train from one or more Rust-generated .npz shards:

cd python/train
uv run python train.py path/to/shard.npz

Export a checkpoint to ONNX:

cd python/train
uv run python export_onnx.py checkpoints/gomoku.pt

Expected shard arrays:

• obs: [N, 3, 15, 15], float32.
• policy: [N, 225], float32, normalized visit distribution.
• value: [N], float32, final outcome from the saved position's current-player perspective.
• mask: [N, 225], bool or uint8.

Core Interfaces

Actions are encoded as u16 values in 0..225:

action = row * 15 + col
row = action / 15
col = action % 15

Observation channels are encoded from the current player's perspective:

• Channel 0: current player's stones.
• Channel 1: opponent's stones.
• Channel 2: constant plane filled with 1.0.

Network outputs expected by the Rust search interface:

• policy_logits: [B, 225].
• value: [B, 1], in [-1, 1], from the current player's perspective.

Development Notes

The Rust core intentionally has no GUI, ndarray, or ONNX dependency yet. That keeps the rules, environment, and search layers fast to compile and easy to test while the AlphaZero pipeline is being built out.