Transformer Chess Engine

Encoder Transformer chess engine made with from scratch only with pytorch, numpy trained on 5M Stockfish games, which repeatedly makes btter moves than 2200elo rated player in familiar positions.

Status: Personal project

The engine demonstrates strong strategic understanding (playing at ~2200 Elo level in familiar positions) but currently suffers from Distribution Shift in endgames and "weird" positions.

We have built the infrastructure to fix this. We are looking for contributors with strong hardware (NVIDIA 3090/4090s) to run the final training phase (see Help Wanted).

---

Performance & The "Super-Engine" Issue

The Good

The model is excellent at:

• Openings: Plays standard book moves with high confidence.
• Midgame Strategy: In standard positions, it finds complex positional ideas.
• Example: In testing against a 2200 Elo rated bot, it achieved equal or winning positions in the middlegame consistently.

The Issue

The model was trained purely on high-level games (Stockfish vs Stockfish, GM vs GM). Problem: It has never seen a blunder.

• When the game drifts into a "weird" position (which never happens in super-GM games), the model hits Out-Of-Distribution (OOD) states.
• It "hallucinates" (thinking it is winning when losing, or vice versa) or plays random moves because it doesn't know how to punish mistakes.
• Deep Search helps, but often the model's evaluation is so broken in these states that search cannot recover.

The Solution (Implemented!)

The fix is Correction Training (also known as Dagger or Expert Iteration):

1. Let the model play and make mistakes.
2. Use Stockfish to label the correct move in those specific "bad" positions.
3. Fine-tune the model on this new dataset.

We have implemented the full infrastructure for this. (See below).

---

Installation & Usage

1. Requirements

pip install torch chess numpy

2. Play Against the Bot

You can play the bot via the CLI. Note: The Search (MCTS-like) is disabled by default because on standard hardware (like an RTX 2070), the OOD hallucinations makes search unreliable. Raw policy play is often more stable.

# Play with Raw Policy (Recommended for now)
python -m src.cli --model checkpoints/best_model.pt --color white --searchOff

# Enable Search (Experimental)
python -m src.cli --model checkpoints/best_model.pt --color white --searchOn

---

Help Wanted: Correction Training

I (@author) am limited by hardware (RTX 2070 Super), which makes generating the necessary correction dataset too slow. If you have a strong GPU, you can finish this project.

We have provided two scripts to solve the OOD issue:

Step 1: Generate Correction Data

This script runs the model against itself. When it reaches a weird position, it asks Stockfish for the truth.

# Generates games and saves (FEN, BestMove) pairs to correction_data.jsonl
python -m src.training.stockfish_generator \
  --model checkpoints/best_model.pt \
  --stockfish "path/to/stockfish" \
  --out correction_data.jsonl \
  --games 1000  # Recommend: 10,000+

Step 2: Fine-Tune

This script fine-tunes the model on the corrected data, teaching it how to verify and recover.

python -m src.training.finetune \
  --model checkpoints/best_model.pt \
  --data correction_data.jsonl \
  --epochs 1 # Keep it low to avoid forgetting openings