Train Your Own Chess AI. Watch your creation defeat you | by Logan Spears | Towards Data Science

TLDR: Colab Link
Deep Learning has taken over the computer chess world. In 2017 AlphaZero, an iteration of AlphaGo targeting chess, mesmerized the chess fans by handlely defeating Stockfish. New developments in the space such as Leela Zero (an open source AlphaZero implementation that happens to use my chess lib) and Stockfish NNUE (efficiently updatable neural net reversed) show neural nets will continue to dominate chess for the foreseeable future.
As a chess enthusiast and AI practitioner, I set out to create my own chess AI but was discouraged by a daunting rumor: AlphaZero cost $35MM to train. AlphaZero trains entirely through reinforcement learning and self play to avoid outside dependencies. While obviously effective, self play is incredibly inefficient from a cost perspective. Today’s engines can already tell the difference between a good position and bad position. This prior knowledge can be used to bootstrap the learning process and make it possible to train a chess AI with very limited cost (or free on Colab). Lets try it! And train a neural net using supervised machine learning with the dataset coming from existing engines.
How does an engine determine if a position is good? Given a position, engines return a score of how good the position is from white perspective. This score is called an evaluation, typically shortened to eval, and the eval function is at the heart of chess engines. The eval is a realitive strength measurement denominated in pawn equivalents. For example if white is up a pawn, the evaluation will be +1. If black is up a knight the eval will be -3. Raw material isn’t the only thing that determines the eval, positional information and future moves are also taken into account.
Traditionally the eval function was implemented by handcrafted algorithms measuring concepts like material imbalance, piece mobility, and king safety quantitatively. These direct measurement, combined with iterative deeping alone produce super human results. Neural networks can substitute for these handcrafted algorithms and return the eval directly eliminating the need for specialized coding. Overall the eval function takes in a position and returns an evaluation score. How do we turn that into a model?
The output, eval, can be represented by a single float based neuron. The input, representing the board position, can be encoded using bitboards (64 bits, one for each square on the chess board) for each piece type and a few remaining bits for move index, color to move, and en passant square. Together this input data forms a string of 808 bits (1s and 0s) that can be converted to floats and inputted into the model directly.
Now that we planned out our inputs, outputs, and general model concept, lets start building a dataset.
I play on Lichess.com because it’s free, open source, and run by developers. Lichess hosts downloadable monthly shards of every game played on the website. July 2021 contains 92+ million games. In the notes section it mentioned a very important detail: “About 6% of the games include Stockfish analysis evaluations” which in this case would be 5.5+ million games. With the average chess game lasting around 40 moves, 80 positions (each side gets a turn per “move”), July should contain 441,600,000 non-unique positions. This seems sufficient for creating our initial dataset.
Lichess’ monthly shards are encoded as .pgn.bz2 files. The file extension .bz2 is a form of compression and .pgn is “Portable Game Notation” which is a text based format for encoding chess games by as a series of moves in Algebraic Notation. Here is the text of a single PGN file record:
[Event "Rated Bullet tournament https://lichess.org/tournament/yc1WW2Ox"][Site "https://lichess.org/PpwPOZMq"][Date "2017.04.01"][Round "-"][White "Abbot"][Black "Costello"][Result "0-1"][UTCDate "2017.04.01"][UTCTime "11:32:01"][WhiteElo "2100"][BlackElo "2000"][WhiteRatingDiff "-4"][BlackRatingDiff "+1"][WhiteTitle "FM"][ECO "B30"][Opening "Sicilian Defense: Old Sicilian"][TimeControl "300+0"][Termination "Time forfeit"]1. e4 { [%eval 0.17] [%clk 0:00:30] } c5 { [%eval 0.19] [%clk 0:00:30] }2. Nf3 { [%eval 0.25] [%clk 0:00:29] } 2... Nc6 { [%eval 0.33] [%clk 0:00:30] }3. Bc4 { [%eval -0.13] [%clk 0:00:28] } 3... e6 { [%eval -0.04] [%clk 0:00:30] }4. c3 { [%eval -0.4] [%clk 0:00:27] } 4... b5? { [%eval 1.18] [%clk 0:00:30] }5. Bb3?! { [%eval 0.21] [%clk 0:00:26] } 5... c4 { [%eval 0.32] [%clk 0:00:29] }6. Bc2 { [%eval 0.2] [%clk 0:00:25] } 6... a5 { [%eval 0.6] [%clk 0:00:29] }7. d4 { [%eval 0.29] [%clk 0:00:23] } 7... cxd3 { [%eval 0.6] [%clk 0:00:27] }8. Qxd3 { [%eval 0.12] [%clk 0:00:22] } 8... Nf6 { [%eval 0.52] [%clk 0:00:26] }9. e5 { [%eval 0.39] [%clk 0:00:21] } 9... Nd5 { [%eval 0.45] [%clk 0:00:25] }10. Bg5?! { [%eval -0.44] [%clk 0:00:18] } 10... Qc7 { [%eval -0.12] [%clk 0:00:23] }11. Nbd2?? { [%eval -3.15] [%clk 0:00:14] } 11... h6 { [%eval -2.99] [%clk 0:00:23] }12. Bh4 { [%eval -3.0] [%clk 0:00:11] } 12... Ba6? { [%eval -0.12] [%clk 0:00:23] }13. b3?? { [%eval -4.14] [%clk 0:00:02] } 13... Nf4? { [%eval -2.73] [%clk 0:00:21] } 0-1
After each move, within a comments section, Stockfish based evaluations are written in a custom text format.
1. e4 { [%eval 0.24] [%clk 0:05:00] }
This gives us eval paired with position (computed from the series of moves). The original compressed monthly shard for July 2021 is 25.2 GB so a single CSV file is likely out of the question. An SQLite database with proper indexes seems more appropriate. The schema includes the FEN (text encoding of the position for debugging), binary for direct use by the model, and the evaluation for the training output:
> select * from evaluations limit 7;id|fen|binary|eval1|rnbqkbnr/pppppppp/8/8/3P4/8/PPP1PPPP/RNBQKBNR b KQkq d3 0 1|0.02|rnbqkbnr/ppp1pppp/8/3p4/3P4/8/PPP1PPPP/RNBQKBNR w KQkq d6 0 2|0.273|rnbqkbnr/ppp1pppp/8/3p4/2PP4/8/PP2PPPP/RNBQKBNR b KQkq c3 0 2|0.04|rnbqkbnr/ppp2ppp/4p3/3p4/2PP4/8/PP2PPPP/RNBQKBNR w KQkq - 0 3|0.095|rnbqkbnr/ppp2ppp/4p3/3P4/3P4/8/PP2PPPP/RNBQKBNR b KQkq - 0 3|0.16|rnbqkbnr/ppp2ppp/8/3p4/3P4/8/PP2PPPP/RNBQKBNR w KQkq - 0 4|0.07|rnbqkbnr/ppp2ppp/8/3p4/3P4/2N5/PP2PPPP/R1BQKBNR b KQkq - 1 4|0.09
Now that we have our data lets start training!
For this project lets use Pytorch Lightning which is a convience layer on top of Pytorch. In order to train, we have to primaily implement two classes: Dataset and LightningModule. Due to the large size of the dataset (the db ended up with 37MM positions), an IterableDataset makes the most sense and treats the data as an infinite stream. The __next__ method just needs to return a random row of the data. Since we aren’t worried about overfitting on tens of millions of positions, we don’t need to split the data into training, validation, and test sets.
EvaluationDataset takes single random row from the SQLite table and preprocesses it by extracting the binary value in raw bytes, converting those bytes to floats using numpy’s frombuffer and unpackbits functions, and forming the required 808 length float array as input. The evaluation value is is extracted and bounded between -15 and 15. Bounding protects against extremely winning position dominating the loss function (such as +45 white). +-15 is very certainly winning for each side and evaluating past that is unnecessary for our purposes and will slow down training convergence.
The model consists of N layers (initially set to 4) of 808 neurons until a final output layer of 1 neural which directly outputs the evaluation. An L1 loss, or the absolute average difference between labels and predictions, is used because the output is intuitive. For example if the L1 loss is 2.0, then on average the model differs from stockfish by 2 pawns of evaluation.
Success, loss is going down! Eventually two versions of our model were trained. One with four layers reaching a loss of 1.29 and one with six layers and a loss of 1.06.
The best way to test our AI is to play it! Using alpha beta search, transposition tables, and few other standard chess AI techniques we can create a full chess engine to play against. The AI client was implemented in Go (so I could use my chess lib) and the model was hosted on a colocated python based gRPC server serving the model on Nvidia 2080 accelerated linux server I have at home.
The Go based chess engine is less than 300 LOC, here is the code for the primary Move method.
And here is the pthon gRPC server implementation serving the model. It takes batches of encoded positions and returns evaluations.
Now that we have our bot hooked up and our model serving board evaluations we are ready to play! The bot’s opponent (myself) has a measly ~1200 average rating on Lichess across rapid, blitz, and bullet. There were no time controls on these games. I took as long as a I needed and the AI took around 20 secs per move. In the two games I played against the bot, I charitably gave myself white.
Medium unfortunately doesn’t allow the embeddings so I have included a Gif of the entire game as well as a link to the full game on Lichess complete with a Stockfish analysis.
lichess.org
I played E4 (I always play E4), and the AI chose C5 known as the “Sicilian Defense”. The first move by black is fascinating because the AI has no opening book and no game play experience, yet chose one of the most popular defenses for black. The opening was pretty balanced with no side gaining a clear advantage until I trade my dark square bishop for the black knight on move 10. In hindsight, this opened up black’s dark square bishop and exposed the weakness on the B2 square. I then compounded the error with another bishop for knight trade on move 13 which dropped a pawn and allowed attacking lines from black’s queen. I hunkered down and defended with the best of my ability and got the game close to even by move 25. By neglecting to take my pawn, black underestimated the potential promotion on E8 and blundered on move 26 allowing knight to B3 with the idea that if the knight takes I can force a rook trade and the bishop can’t defend the A pawn from promotion. Unfortunately I didn’t see that either and I let my temporary advantage slip away. The game ended with a serious of uneventful moves and I agreed to a draw (with Stockfish agreeing the position was even).
lichess.org
Opening with the French, black allowed me to take the center by move 8 resulting in a position stockfish evaluates +1.1 for white. Interestingly Stockfish wanted me to castle queen side and push my advantage in the center. I have a habit of castling king side and played in a more safe style to avoid unfamiliar play. On move 11, black offers a queen trade that I accepted to take the queens off the board reducing tactical opportunities that could favor the AI. I was doing well until I blundered on move 16. I didn’t see that black could skewer my rooks with its light square bishop and win an exchange. With a pawn push on move 32 not resulting in winning chances (even with the balancing of rooks), black gradually picked off my pawns forcing me to resign.
While the AI won, I couldn’t tell if the computer had gotten better or my blunder made the comparison unfair. The fairest way to settle the matter was have the AI versions play each other.
lichess.org
Initially I gave white to the 4 layer AI, but it ended in a bland draw so I revered the color assignments. The second game resulted in an incredibly exciting match. The fast paced action started on move 12 with the bishop taking the G7 pawn and forking the knight and rook on H6 and H8 respectively. Black responded with a wild sacrifice using the dark square bishop followed by a queen check. White calmly deflected and finally claimed the hanging rook. On move 16 black sacrificed another piece, this time a knight, in order to keep the attack on white’s king alive. After defending and throwing in a queen check of its own, white confidently checked the black king on move 22. This was especially dangerous because black, if given a free move, can checkmate in one! White found the force mating sequence, prevailed in an exhilarating checkmate, and proved that its lower training loss resulted in better chess.
The most remarkable aspect of analyzing these games is how “human” the AI feels. Common openings, waiting for blunders, and attempting comebacks with wild sacrifices are standard fair in online human play.
There is certainly room for improvement in our approach. A lot of well known techniques were not used in order to keep this project’s complexity manageable and focus on the AI portions:
Move ordering to prioritize searching “interesting” move sequences
Endgame Tablebase for precomputed winning endgame positions
Opening Book for precomputed openings
Only search a couple hundred thousand positions at 4 ply
Many others
There is also some flaws in our simplistic model and training procedures
Move counts are inputted into the model in bits instead of the number as a float
The model architecture was the first design attempted with no optimization
Learning rate and batch size schedulers could result in lower total loss
Chess positions are unique in the database and therefore later positions that are more likely to be unique are over represented in the dataset
Most games are played by amateurs and the search space of superior play might be under represented
The model could be trained for MUCH longer
Overall this was a lot of fun! Supervised training on existing engine evaluation is a faster and more economical path to weight initialization than self play and monte carlo tree search (a la AlphaZero) and a playable model can be generated in a few GPU hours. Hope you enjoyed this post and I look forward to a link of your AI’s games in the comments!