From 5ccfbeb4e88f6c9ffba51331636c96318100b960 Mon Sep 17 00:00:00 2001 From: bbig Date: Thu, 1 Jun 2023 19:00:56 +0800 Subject: [PATCH] some try for policy_gradient --- agents/pg1.py | 95 ++++++++++++++++++++++++ agents/pg2.py | 124 ++++++++++++++++++++++++++++++++ agents/policy_gradient_agent.py | 106 +++++++++++++++++---------- alphago_policy_rl.py | 6 +- encoders/alphago.py | 4 ++ 5 files changed, 295 insertions(+), 40 deletions(-) create mode 100644 agents/pg1.py create mode 100644 agents/pg2.py diff --git a/agents/pg1.py b/agents/pg1.py new file mode 100644 index 0000000..9a3608b --- /dev/null +++ b/agents/pg1.py @@ -0,0 +1,95 @@ +"""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) diff --git a/agents/pg2.py b/agents/pg2.py new file mode 100644 index 0000000..5d885b9 --- /dev/null +++ b/agents/pg2.py @@ -0,0 +1,124 @@ +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) diff --git a/agents/policy_gradient_agent.py b/agents/policy_gradient_agent.py index e580fc8..5d885b9 100644 --- a/agents/policy_gradient_agent.py +++ b/agents/policy_gradient_agent.py @@ -1,9 +1,8 @@ -"""Policy gradient learning.""" -import numpy as np import torch import torch.nn.functional as F -from torch.optim import SGD +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 @@ -11,15 +10,19 @@ 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) + 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): - """An agent that uses a deep policy network to select moves.""" def __init__(self, model, encoder): - super().__init__() + Agent.__init__(self) self._model = model self._encoder = encoder self._collector = None @@ -27,11 +30,8 @@ class PolicyAgent(Agent): 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] + input_tensor = torch.tensor([encoded_state]) + return self._model(input_tensor)[0] def set_temperature(self, temperature): self._temperature = temperature @@ -42,28 +42,37 @@ class PolicyAgent(Agent): def select_move(self, game_state): num_moves = self._encoder.board_width * self._encoder.board_height - move_probs = self.predict(game_state) + board_tensor = self._encoder.encode(game_state) + x = torch.tensor([board_tensor]) - 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) + if torch.rand(1).item() < self._temperature: + move_probs = torch.ones(num_moves) / num_moves + else: + move_probs = self._model(x)[0] - candidates = np.arange(num_moves) - ranked_moves = np.random.choice(candidates, num_moves, replace=False, p=move_probs) + 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) + 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): + 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) + 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=1e-7, clipnorm=1.0, batch_size=512): - opt = SGD(self._model.parameters(), lr=lr) + 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)) @@ -72,23 +81,44 @@ class PolicyAgent(Agent): 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() + 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(self._model.state_dict(), 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, encoder): + 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)) + # 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) diff --git a/alphago_policy_rl.py b/alphago_policy_rl.py index a0e10dc..4e71199 100644 --- a/alphago_policy_rl.py +++ b/alphago_policy_rl.py @@ -19,12 +19,14 @@ alphago_rl_agent = PolicyAgent(sl_agent.model.to(device), encoder) opponent = PolicyAgent(sl_opponent.model.to(device), encoder) # Run simulation -num_games = 1000 +# num_games = 1000 +num_games = 10 +batch_size = 1024 with tqdm(total=num_games) as pbar: experience = experience_simulation(num_games, alphago_rl_agent, opponent, progress_callback=lambda: pbar.update(1)) -alphago_rl_agent.train(experience) +alphago_rl_agent.train(experience, batch_size=batch_size) # # Serialize RL agent # alphago_rl_agent.save('checkpoints/alphago_rl_policy.pt') diff --git a/encoders/alphago.py b/encoders/alphago.py index 4c9f7dc..ff0787e 100644 --- a/encoders/alphago.py +++ b/encoders/alphago.py @@ -133,3 +133,7 @@ class AlphaGoEncoder(Encoder): def shape(self): return self.num_planes, self.board_height, self.board_width + + +def create(board_size): + return AlphaGoEncoder(board_size) \ No newline at end of file