125 lines
4.5 KiB
Python
125 lines
4.5 KiB
Python
import torch
|
|
import torch.nn.functional as F
|
|
from torch import optim
|
|
|
|
from tugo import encoders
|
|
from tugo.agents.base import Agent
|
|
from tugo.agents.helpers import is_point_an_eye
|
|
# from tugo.game_logic import goboard
|
|
from tugo.game_logic import goboard_fast as goboard
|
|
|
|
|
|
def policy_gradient_loss(y_true, y_pred):
|
|
clip_pred = torch.clamp(y_pred, torch.finfo(torch.float32).eps, 1 - torch.finfo(torch.float32).eps)
|
|
loss = -1 * y_true * torch.log(clip_pred)
|
|
return torch.mean(torch.sum(loss, dim=1))
|
|
|
|
|
|
def normalize(x):
|
|
total = torch.sum(x)
|
|
return x / total
|
|
|
|
|
|
class PolicyAgent(Agent):
|
|
def __init__(self, model, encoder):
|
|
Agent.__init__(self)
|
|
self._model = model
|
|
self._encoder = encoder
|
|
self._collector = None
|
|
self._temperature = 0.0
|
|
|
|
def predict(self, game_state):
|
|
encoded_state = self._encoder.encode(game_state)
|
|
input_tensor = torch.tensor([encoded_state])
|
|
return self._model(input_tensor)[0]
|
|
|
|
def set_temperature(self, temperature):
|
|
self._temperature = temperature
|
|
|
|
def set_collector(self, collector):
|
|
self._collector = collector
|
|
|
|
def select_move(self, game_state):
|
|
num_moves = self._encoder.board_width * self._encoder.board_height
|
|
|
|
board_tensor = self._encoder.encode(game_state)
|
|
x = torch.tensor([board_tensor])
|
|
|
|
if torch.rand(1).item() < self._temperature:
|
|
move_probs = torch.ones(num_moves) / num_moves
|
|
else:
|
|
move_probs = self._model(x)[0]
|
|
|
|
eps = 1e-5
|
|
move_probs = torch.clamp(move_probs, eps, 1 - eps)
|
|
move_probs = move_probs / torch.sum(move_probs)
|
|
|
|
candidates = torch.arange(num_moves)
|
|
ranked_moves = torch.multinomial(move_probs, num_moves, replacement=False)
|
|
for point_idx in ranked_moves:
|
|
point = self._encoder.decode_point_index(point_idx.item())
|
|
if game_state.is_valid_move(goboard.Move.play(point)) and \
|
|
not is_point_an_eye(game_state.board,
|
|
point,
|
|
game_state.next_player):
|
|
if self._collector is not None:
|
|
self._collector.record_decision(
|
|
state=board_tensor,
|
|
action=point_idx.item()
|
|
)
|
|
return goboard.Move.play(point)
|
|
return goboard.Move.pass_turn()
|
|
|
|
def train(self, experience, lr=0.0000001, clipnorm=1.0, batch_size=512):
|
|
opt = optim.SGD(self._model.parameters(), lr=lr, clipnorm=clipnorm)
|
|
|
|
n = experience.states.shape[0]
|
|
num_moves = self._encoder.board_width * self._encoder.board_height
|
|
y = torch.zeros((n, num_moves))
|
|
for i in range(n):
|
|
action = experience.actions[i]
|
|
reward = experience.rewards[i]
|
|
y[i][action] = reward
|
|
|
|
self._model.train()
|
|
opt.zero_grad()
|
|
outputs = self._model(experience.states)
|
|
loss = F.cross_entropy(outputs, y)
|
|
loss.backward()
|
|
opt.step()
|
|
|
|
def serialize(self, h5file):
|
|
h5file.create_group('encoder')
|
|
h5file['encoder'].attrs['name'] = self._encoder.name()
|
|
h5file['encoder'].attrs['board_width'] = self._encoder.board_width
|
|
h5file['encoder'].attrs['board_height'] = self._encoder.board_height
|
|
h5file.create_group('model')
|
|
torchutil.save_model_to_hdf5_group(self._model, h5file['model'])
|
|
|
|
def save(self, file_path):
|
|
torch.save({
|
|
'model_state_dict': self._model.state_dict(),
|
|
'encoder_name': self._encoder.name(),
|
|
'encoder_board_width': self._encoder.board_width,
|
|
'encoder_board_height': self._encoder.board_height,
|
|
'temperature': self._temperature,
|
|
}, file_path)
|
|
|
|
@classmethod
|
|
def load(cls, file_path):
|
|
device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
|
# model = AlphaGoModel(encoder.get_input_shape(), is_policy_net=True).to(device)
|
|
# model.load_state_dict(torch.load(file_path))
|
|
# return cls(model, encoder)
|
|
|
|
checkpoint = torch.load(file_path)
|
|
model = TheModelClass() # TheModelClass should be replaced with your model class
|
|
model.load_state_dict(checkpoint['model_state_dict'])
|
|
encoder_name = checkpoint['encoder_name']
|
|
board_width = checkpoint['encoder_board_width']
|
|
board_height = checkpoint['encoder_board_height']
|
|
encoder = encoders.get_encoder_by_name(
|
|
encoder_name,
|
|
(board_width, board_height))
|
|
return cls(model, encoder)
|