ReinforcementRenju

A GPU-accelerated AlphaZero implementation for Renju (Gomoku/Five-in-a-Row) using JAX/Flax with a real-time web-based game interface.

Features

• GPU-Accelerated Training: Fully GPU-based MCTS using JAX and Google DeepMind's mctx library
• AlphaZero Algorithm: Self-play reinforcement learning with neural network policy and value heads
• Real-time Web Interface: Human vs AI gameplay via WebSocket with live AI analysis
• Win Rate Visualization: Real-time win probability display and game history tracking
• Renju Rules: Standard 15x15 board implementation

Project Structure

ReinforcementRenju/
├── model/                  # Core ML training modules
│   ├── game.py            # JAX-based game logic (15x15 board)
│   ├── mcts.py            # GPU-accelerated MCTS using mctx
│   ├── network.py         # Flax ResNet architecture
│   ├── train.py           # AlphaZero training loop
│   └── visualization/     # Training metrics and visualization
├── client/
│   ├── backend/           # FastAPI WebSocket server
│   │   ├── server.py      # Real-time game server
│   │   ├── player.py      # Human & AI player classes
│   │   └── game_manager.py
│   └── frontend/          # Vue.js 3 web interface
├── playground/            # Legacy NumPy-based implementation
├── checkpoints_jax_small/ # Pre-trained model checkpoint
├── train_jax.py          # Training entry point
└── run_server.py         # Server launcher

Requirements

Training

• Python 3.8+
• JAX with CUDA support
• Flax, Optax, mctx
• NumPy

Game Server

• FastAPI
• Uvicorn

Frontend

• Node.js 18+
• Vue.js 3
• Vite 5

Installation

1. Python Dependencies

pip install jax jaxlib flax optax mctx fastapi uvicorn numpy

For GPU support (CUDA 11.8):

pip install jax[cuda11_pip] -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

2. Frontend Dependencies

cd client/frontend
npm install

Usage

Training

# Full training (A100 configuration)
python train_jax.py

# Small model training (for testing/development)
python train_jax.py --small

# Resume from checkpoint
python train_jax.py --resume checkpoints_jax/checkpoint_50.pkl

Training Configurations:

Config	Channels	Blocks	Games/Iter	Batch Size	Target GPU
A100	256	20	2048	1024	A100 80GB
Small	128	6	256	128	Consumer

Running the Game Server

Start Backend:

python run_server.py

Start Frontend:

cd client/frontend
npm run dev

Open http://localhost:5173 in your browser to play against the AI.

Technical Details

Neural Network Architecture

• ResNet-style architecture with configurable residual blocks
• Dual output heads:
  • Policy head: Move probability distribution (225 positions)
  • Value head: Win probability estimation (-1 to 1)

MCTS Implementation

• GPU-accelerated tree search using mctx library
• Parallel game simulation for efficient self-play
• Dirichlet noise injection for exploration
• Temperature-based move selection

Game Features

• Standard 15x15 Renju/Gomoku board
• Win condition: 5 stones in a row (horizontal, vertical, diagonal)
• Real-time AI analysis with win rate display
• Move history visualization

Pre-trained Model

A pre-trained small model is included in checkpoints_jax_small/best.pkl. This model was trained using the small configuration and can be used directly for playing against the AI.

Author

Hoyeon Cho, Korea Maritime and Ocean University (KMOU)

License

This project is licensed under the PolyForm Noncommercial License 1.0.0.

Key points:

• Free for research, education, and non-commercial use
• Modifications allowed with proper attribution
• Not a copyleft license (your modifications don't have to use the same license)
• Commercial use requires separate licensing agreement

See the LICENSE file for full terms.