Spaces:
Running
Running
| import numpy as np | |
| import chess | |
| import torch | |
| from o2_model import O2Net, board_to_tensor | |
| from mcts import MCTS | |
| import random | |
| # Optional: Endgame tablebase and opening book integration placeholders | |
| # You can use python-chess's tablebase and opening book modules if desired | |
| # Example for endgame tablebase: | |
| # from chess import tablebase | |
| # tb = tablebase.Tablebase() | |
| # tb.add_tablebase('/path/to/syzygy') | |
| # if tb.probe_wdl(board) is not None: | |
| # # Use tablebase move | |
| # Example for opening book: | |
| # from chess.polyglot import open_reader | |
| # with open_reader('book.bin') as reader: | |
| # entry = reader.find(board) | |
| # move = entry.move | |
| class O2Agent: | |
| def __init__(self, model_path=None): | |
| self.model = O2Net() | |
| if model_path: | |
| self.model.load_state_dict(torch.load(model_path)) | |
| self.model.eval() | |
| def select_move(self, board, use_mcts=True, simulations=100, temperature=0.0): | |
| if use_mcts: | |
| mcts = MCTS(self.model, simulations=simulations) | |
| return mcts.run(board, temperature=temperature) | |
| tensor = torch.tensor(board_to_tensor(board)).unsqueeze(0) | |
| with torch.no_grad(): | |
| policy, _ = self.model(tensor) | |
| legal_moves = list(board.legal_moves) | |
| move_scores = [] | |
| for move in legal_moves: | |
| move_idx = self.move_to_index(move) | |
| move_scores.append(policy[0, move_idx].item()) | |
| if temperature and temperature > 0: | |
| # Softmax sampling | |
| scores = np.array(move_scores) | |
| exp_scores = np.exp(scores / temperature) | |
| probs = exp_scores / np.sum(exp_scores) | |
| move = np.random.choice(legal_moves, p=probs) | |
| return move | |
| best_move = legal_moves[int(torch.tensor(move_scores).argmax())] | |
| return best_move | |
| def move_to_index(self, move): | |
| # Encode move as from_square * 64 + to_square + promotion_offset | |
| from_square = move.from_square | |
| to_square = move.to_square | |
| promotion = move.promotion if move.promotion else 0 | |
| promotion_offset = 0 | |
| if promotion: | |
| # Promotion: 1=Knight, 2=Bishop, 3=Rook, 4=Queen (python-chess) | |
| # Offset: 4096 + (promotion-1)*64*64//4 | |
| promotion_offset = 4096 + (promotion - 1) * 256 | |
| idx = from_square * 64 + to_square + promotion_offset | |
| # Ensure index is within bounds | |
| return idx if idx < 4672 else idx % 4672 | |
| def index_to_move(self, board, index): | |
| # Decode index to move (reverse of move_to_index) | |
| if index >= 4096: | |
| promotion = (index - 4096) % 4 + 1 | |
| idx = index - 4096 | |
| from_square = idx // 64 | |
| to_square = idx % 64 | |
| move = chess.Move(from_square, to_square, promotion=promotion) | |
| else: | |
| from_square = index // 64 | |
| to_square = index % 64 | |
| move = chess.Move(from_square, to_square) | |
| if move in board.legal_moves: | |
| return move | |
| # Fallback: pick a random legal move | |
| return random.choice(list(board.legal_moves)) | |
| if __name__ == "__main__": | |
| board = chess.Board() | |
| agent = O2Agent() | |
| move = agent.select_move(board) | |
| print("O2 selects:", move) | |