Update chess-gpt-eval-contrastive/main.py

chess-gpt-eval-contrastive/main.py

@@ -395,6 +395,49 @@ def add_random_moves(
     return game_state, board, num_moves
 
 
+def evaluate_position(fen, backend):
+    gamestate = GameState(fen=fen)
+    result = backend.evaluate(gamestate.as_input(backend))[0]
+    return result.q()
+
+
+def material_balance(board):
+    PV = {
+        'pawn': 1,
+        'knight': 3,
+        'bishop': 3,
+        'rook': 5,
+        'queen': 9,
+        'king': 0
+    }
+    
+    if board.is_insufficient_material():
+        return 0
+
+    wp = len(board.pieces(chess.PAWN, chess.WHITE))
+    bp = len(board.pieces(chess.PAWN, chess.BLACK))
+
+    wn = len(board.pieces(chess.KNIGHT, chess.WHITE))
+    bn = len(board.pieces(chess.KNIGHT, chess.BLACK))
+
+    wb = len(board.pieces(chess.BISHOP, chess.WHITE))
+    bb = len(board.pieces(chess.BISHOP, chess.BLACK))
+
+    wr = len(board.pieces(chess.ROOK, chess.WHITE))
+    br = len(board.pieces(chess.ROOK, chess.BLACK))
+
+    wq = len(board.pieces(chess.QUEEN, chess.WHITE))
+    bq = len(board.pieces(chess.QUEEN, chess.BLACK))
+
+    return (
+        PV['pawn'] * (wp - bp) +
+        PV['knight'] * (wn - bn) +
+        PV['bishop'] * (wb - bb) +
+        PV['rook'] * (wr - br) +
+        PV['queen'] * (wq - bq)
+    )
+
+
 # Return is (move_san, move_uci, attempts, is_resignation, is_illegal_move)
 def get_legal_move(
     player: Player,
@@ -554,6 +597,8 @@ def play_games(
             #print(f"|{game_state}|")
             #print(f"{current_move_num}", end=" ")
 
+            if update_linear:
+                prev_q_value = evaluate_position(board.fen(), player_two.backend)
             (
                 game_state,
                 player_one_resignation,
@@ -574,6 +619,12 @@ def play_games(
             ):
                 break
 
+            if update_linear:
+                curr_q_value = evaluate_position(board.fen(), player_two.backend)
+                q_value_delta = curr_q_value - prev_q_value
+                material_bal = material_balance(board)
+                player_one.update_linear_probe_targets(curr_q_value, q_value_delta, material_bal)
+            
             (
                 game_state,
                 player_two_resignation,
@@ -639,6 +690,9 @@ def play_games(
         
                 if games_saved % save_activations_every == 0:
                     player_one.save_activations(activations_path)
+
+            if update_linear and games_saved % save_activations_every == 0:
+                player_one.save_linear_probe_data(linear_path)
         else:
             print("Duplicate game; not saved.")
             
@@ -675,11 +729,16 @@ move_num_in_gamestate = False
 book_opening = False
 random_opening = True
 random_opening_moves = 10
+
 activations_path="activations_rdm.pkl"
-update_activations = True
+update_activations = False
 apply_activations = False
 save_activations_every = 25
 contrastive_weight = 0.8
+
+linear_path="linear.pkl"
+update_linear = True
+eval_linear = False
 if __name__ == "__main__":
     for nanogpt_player in player_ones:
         i = 0