52 lines
2.0 KiB
Python
52 lines
2.0 KiB
Python
from tugo.agents.base import Agent
|
|
from tugo.agents.helpers import is_point_an_eye
|
|
# from game_logic import goboard
|
|
from tugo.game_logic import goboard_fast as goboard
|
|
from tugo.models import AlphaGoModel
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
class DeepLearningAgent(Agent):
|
|
def __init__(self, model, encoder):
|
|
super().__init__()
|
|
self.model = model
|
|
self.encoder = encoder
|
|
|
|
def predict(self, game_state):
|
|
encoded_state = self.encoder.encode(game_state)
|
|
input_tensor = torch.tensor([encoded_state], dtype=torch.float32).to('cuda')
|
|
with torch.no_grad():
|
|
output_tensor = self.model(input_tensor)
|
|
return output_tensor.cpu().numpy()[0]
|
|
|
|
def select_move(self, game_state):
|
|
num_moves = self.encoder.board_width * self.encoder.board_height
|
|
move_probs = self.predict(game_state)
|
|
|
|
move_probs = move_probs ** 3
|
|
eps = 1e-6
|
|
move_probs = np.clip(move_probs, eps, 1 - eps)
|
|
move_probs = move_probs / np.sum(move_probs)
|
|
|
|
candidates = np.arange(num_moves)
|
|
ranked_moves = np.random.choice(
|
|
candidates, num_moves, replace=False, p=move_probs)
|
|
for point_idx in ranked_moves:
|
|
point = self.encoder.decode_point_index(point_idx)
|
|
if game_state.is_valid_move(goboard.Move.play(point)) and \
|
|
not is_point_an_eye(game_state.board, point, game_state.next_player):
|
|
return goboard.Move.play(point)
|
|
return goboard.Move.pass_turn()
|
|
|
|
def save(self, file_path):
|
|
torch.save(self.model.state_dict(), file_path)
|
|
|
|
@classmethod
|
|
def load(cls, file_path, encoder):
|
|
device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
|
# model = AlphaGoModel(encoder.get_input_shape(), is_policy_net=True).to(device)
|
|
model = AlphaGoModel(encoder.shape(), is_policy_net=True, num_classes=361).to(device)
|
|
model.load_state_dict(torch.load(file_path))
|
|
return cls(model, encoder)
|