clean 1/n

Author: k-rt-k

CPP_backend.cpp

@@ -1,3 +1,8 @@
+/*
+    MCST implementation in C++ for chess
+    Interfaced with python using boost python
+    Check README for installation and compilation instructions
+*/
 #include <boost/python.hpp>
 #include <Python.h>
 #include <boost/multiprecision/cpp_dec_float.hpp>
@@ -13,7 +18,8 @@
 #include <unordered_map>
 long long int cnt = 0;
 
-#define ll long long int
+typedef long long int ll;
+
 #define AMOUNT_OF_PLANES 73
 #define BOARD_SIZE 8
 #define DIRICHLET_NOISE 0.3
@@ -25,6 +31,7 @@ class C_Edge;
 class C_MCTS;
 class C_Edge;
 
+// One Node in the MCTS tree
 class C_Node{
     public:
         std::string state;
@@ -44,6 +51,7 @@ class C_Node{
         uint64_t get_edge(boost::python::object action);
 };
 
+// One Edge in the MCTS tree, which corresponds to one action
 class C_Edge{   
     public:
         C_Node* input_node;
@@ -63,9 +71,6 @@ class C_Edge{
         uint64_t get_N();
 };
 
-// class C_Action{
-
-// }
 
 class C_MCTS{
     public:
@@ -94,7 +99,7 @@ class C_MCTS{
         ~C_MCTS();
 };     
 
-
+// Recursively delete the MCTS tree beginning at node
 void delete_mcts_tree(C_Node* node){
 
     if(! node) return;
@@ -131,11 +136,9 @@ std::string C_Node::step(boost::python::object action){
 }
 
 bool C_Node::is_game_over(){
-    // Py_Initialize();
     boost::python::object chess_module = boost::python::import("chess");
     boost::python::object board = chess_module.attr("Board")(this->state);
     bool is_game_over = boost::python::extract<bool>(board.attr("is_game_over")());
-    // Py_Finalize();
     return is_game_over;
 }
 
@@ -314,8 +317,7 @@ void C_MCTS::map_valid_move(boost::python::object move){
 
 }
 
-// TODO : Need suggestions of how to implement probabilities to actions
-
+// 
 std::unordered_map<std::string, long double> C_MCTS::probabilities_to_actions(boost::python::object probabilities, std::string bord){
 	std::unordered_map <std::string, long double> actions;
 	boost::python::object chess = boost::python::import("chess");
@@ -484,6 +486,9 @@ std::string C_MCTS::get_edge_uci(uint64_t edge){
     return boost::python::extract<std::string>(((C_Edge*)edge)->action.attr("uci")());
 }
 
+// Python interface for the MCTS class
+// The Names of the objects are the string arguments
+
 BOOST_PYTHON_MODULE(CPP_backend)
 {
     Py_Initialize();

agent.py

@@ -1,18 +1,14 @@
-import numpy as np
+# Contains the Agent class, which is used to play chess moves on the environment.
 import torch
-import torch.nn as nn
-import torch.nn.functional as F
 from neural_network import AgentNetwork
-# from mcts import MCTS # missing
-from CPP_backend import *
-import utils
-import chess
-import time
+from CPP_backend import MCTS
+
+from chess import STARTING_FEN
 import config
 import datetime
 
 class Agent:
-    def __init__(self,local_preds:bool = False, model_path:str|None = None,state:str = chess.STARTING_FEN, device=None):
+    def __init__(self,local_preds:bool = False, model_path:str|None = None,state:str = STARTING_FEN, device=None)->None:
         """
         An agent is an object that can play chessmoves on the environment.
         Based on the parameters, it can play with a local model, or send its input to a server.
@@ -32,14 +28,13 @@ def __init__(self,local_preds:bool = False, model_path:str|None = None,state:str
             if model_path is not None:
                 self.model.load_state_dict(torch.load(model_path))
         else :
-            raise NotImplementedError("Server predictions not implemented yet")
+            raise NotImplementedError("Server predictions not implemented")
 
         self.state = state
         self.mcts = MCTS(self, state, True)
 
-    def run_simulations(self,n:int=1):
+    def run_simulations(self,n:int=1)->None:
         with torch.no_grad():
-
             self.mcts.run_simulations(n)
 
     def save_model(self,timestamped:bool = False)->str:
@@ -53,24 +48,16 @@ def save_model(self,timestamped:bool = False)->str:
 
         return model_path
 
-    def predict(self, data:torch.Tensor):
+    def predict(self, data:torch.Tensor)->torch.Tensor:
         data = torch.Tensor(data).to(torch.float32).unsqueeze(0).to(self.device)
-        # print(data.shape)
-        # print("in agent predict")
         if self.local_preds:
-            # print('local')
             return self.predict_local(data)
         return self.predict_server(data)
 
-    def predict_local(self,data:torch.Tensor):
-        # self.model.eval()
-        
+    def predict_local(self,data:torch.Tensor)->(torch.Tensor,float):
         with torch.no_grad():
             v, p = self.model(data)
             return p.cpu(), v.cpu().item()
 
     def predict_server(self,data:torch.Tensor):
-        raise NotImplementedError("Server predictions not implemented yet")
-  
-if __name__=="__main__":
-    pass
+        raise NotImplementedError("Server predictions not implemented")

chessEnv.py

@@ -1,81 +1,4 @@
-# ### This is the Chess Environment that will interact with the Agent
-
-# ### TODO: make this as c as possible
-
-# import numpy as np
-# import chess
-# import torch
-# from copy import deepcopy
-
-# btoi = lambda x: (1 if x else -1) 
-# strenc = np.array(['r', 'n', 'b', 'q', 'k','p','P','R', 'N', 'B', 'Q', 'K']).reshape(1,12,1,1)
-
-# ### TODO:docs
-# class ChessEnv:
-#     def __init__(self,board:str|list,batch_size:int,board_size:int=8,torch_device:str='cuda'):
-#         self.batch_size = batch_size
-#         self.board_size = board_size
-#         self.torch_device = torch_device
-#         self.num_piecetype = 12
-        
-#         self.__init_board_frm_str(board)
-        
-        
-    
-#     ### Interface functions ###    
-#     def get_embedding(self)->torch.Tensor:
-#         self.__update_embedding()
-#         ## expand out (_,7) tensor to (_,7,8,8) tensor
-#         self.board_states_embedding = self.board_states.unsqueeze(2).unsqueeze(3).repeat(1,1,self.board_size,self.board_size)
-#         return torch.cat([self.board_embedding,self.board_states_embedding],dim=1)
-        
-    
-#     ### Functions to convert between different representations of the board ###
-#     def __init_board_frm_str(self,board:str|list)->None:
-#         if type(board) == str:
-#             self.board_init = (chess.Board(board),)
-#         elif type(board) == list:
-#             self.board_init = [chess.Board(b) for b in board]
-#         self.board = self.board_init.deepcopy()
-        
-#         ### this will be changed !!
-#         self.movenum = 0
-#         ### load the initial embedding ### -> assuming initially history is repeated rather than empty
-#         self.board_embedding = self.__board_to_tensor(self.board).repeat(1,self.board_size,1,1)
-        
-#         self.reps = torch.zeros((self.batch_size,1),device=self.torch_device)
-        
-        
-        
-#     ## remember devices
-    
-#     ## (turn, 4 castling rights,movenum)
-#     def __get_board_states_single(self,board:chess.Board)->None:
-#         return torch.tensor([[btoi(board.turn),btoi(board.has_kingside_castling_rights(chess.WHITE)),btoi(board.has_kingside_castling_rights(chess.BLACK)),btoi(board.has_queenside_castling_rights(chess.WHITE)),btoi(board.has_queenside_castling_rights(chess.BLACK)),self.movenum]],device=self.torch_device)
-#     ## add reps too
-#     def __get_board_states(self)->torch.Tensor:
-#         sixvars =  torch.cat([self.__get_board_states_single(b) for b in self.board],dim=0)
-#         return torch.cat((sixvars,self.reps),dim=1)
-#     # f me #
-#     def __board_to_tensor(self,boards:list|tuple)->torch.Tensor:
-#         arr = (np.array([b.__str__().split() for b in boards]).reshape(-1,1,self.board_size,self.board_size)==strenc)*1
-#         return torch.Tensor(arr, device=self.torch_device)
-        
-#     def __update_embedding(self)->None: ### maybe this works correctly
-#         self.board_embedding = torch.cat([self.board_embedding[:,self.num_piecetype:,:,:],self.__board_to_tensor(self.board)],dim=0)
-#         self.reps = (self.reps + 1)*torch.all(torch.all(torch.all(self.board_embedding[:,-2*self.num_piecetype:-self.num_piecetype,:,:]==self.board_embedding[:,-self.num_piecetype:,:,:],dim=3),dim=2),dim=1,keepdim=True)
-#         self.board_states = self.__get_board_states()
-#         return 
-    
-#     ## convert moves to mask
-#     def __moves_to_mask(self,moves)->torch.Tensor:
-#         pass
-    
-#     ## taking the one-hot encoding of the move chosen and updates the board based on it
-#     def __make_move(self,movetensor)->list:
-#         pass
-    
-    
+### This is the Chess Environment that will interact with the Agent and the MCTS
 #---------------------------- A simpler implementation for now ---------------
 
 import config
@@ -139,11 +62,12 @@ def state_to_input(fen: str):
     chess.KING: 0
 }
 
+### ensure stockfish is installed and the path is correct
 stockfish = Stockfish(os.path.expanduser(config.STOCKFISH))
 
 def estimate_winner(board: chess.Board) -> int:
     """
-    Estimate the winner of the current node.
+    Estimate the winner of the current node. A piece counting heuristic is used.
     Pawn = 1, Bishop = 3, Rook = 5, Queen = 9
     Positive score = white wins, negative score = black wins
     """

collect.py

@@ -1,3 +1,4 @@
+# Description: This file is used to collect experience data from the lichess database.
 import pandas as pd
 import config
 import threading
@@ -12,9 +13,12 @@
 import os
 import numpy as np
 
+# Location of the puzzles.csv file. Ensure the file is present at this location
 CSV_FILE = "puzzles.csv"
+# Number of games to play at once
 N = 5
 
+# Play a puzzle starting from the given fen and moves. This is still self-play
 def play_puzzle(fen, moves):
     model_path = None if len(os.listdir(config.BEST_MODEL)) == 0 else f"{config.BEST_MODEL}best-model.pth"
     white = Agent(local_preds=True, model_path=model_path)
@@ -27,6 +31,7 @@ def play_puzzle(fen, moves):
 
     game.game()
 
+# Play a full game from the start
 def play_normal():
     model_path = None if len(os.listdir(config.BEST_MODEL)) == 0 else f"{config.BEST_MODEL}best-model.pth"
     white=Agent(local_preds=True, model_path=model_path)

config.py

@@ -1,3 +1,5 @@
+# Parameters for all the files #
+
 BOARD_SIZE = 8
 MAX_MOVES = 50
 PREVIOUS_MOVES = 8
@@ -21,6 +23,7 @@
 MEMORY = "./memory/"
 PUZZLE = "./puzzles/"
 BEST_MODEL = "./best_model/"
+PGN = "./pgn/"
 
 #----------Executable Locations--------------------
 STOCKFISH = "~/stockfish/stockfish-ubuntu-x86-64-modern"