96 lines
3.5 KiB
Python
96 lines
3.5 KiB
Python
"""Policy gradient learning."""
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from torch.optim import SGD
|
|
|
|
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, 1e-10, 1 - 1e-10)
|
|
loss = -1 * y_true * torch.log(clip_pred)
|
|
return torch.mean(torch.sum(loss, dim=1))
|
|
|
|
|
|
class PolicyAgent(Agent):
|
|
"""An agent that uses a deep policy network to select moves."""
|
|
def __init__(self, model, encoder):
|
|
super().__init__()
|
|
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], dtype=torch.float32).to('cuda')
|
|
input_tensor = encoded_state.float().to('cuda')
|
|
with torch.no_grad():
|
|
output_tensor = self._model(input_tensor)
|
|
return output_tensor.cpu().numpy()[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)
|
|
|
|
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):
|
|
if self._collector is not None:
|
|
self._collector.record_decision(state=board_tensor, action=point_idx)
|
|
return goboard.Move.play(point)
|
|
|
|
return goboard.Move.pass_turn()
|
|
|
|
def train(self, experience, lr=1e-7, clipnorm=1.0, batch_size=512):
|
|
opt = SGD(self._model.parameters(), lr=lr)
|
|
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
|
|
|
|
for epoch in range(1):
|
|
permutation = torch.randperm(n)
|
|
for i in range(0, n, batch_size):
|
|
indices = permutation[i:i+batch_size]
|
|
batch_x, batch_y = experience.states[indices], y[indices]
|
|
opt.zero_grad()
|
|
outputs = self._model(batch_x)
|
|
loss = F.cross_entropy(outputs, batch_y)
|
|
loss.backward()
|
|
opt.step()
|
|
|
|
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.load_state_dict(torch.load(file_path))
|
|
return cls(model, encoder)
|