File size: 5,237 Bytes
afab5da | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 | import os
import random
import copy
import numpy as np
import torch
import torch.nn.functional as F
from tqdm import trange
from Alpha_MCTS import Alpha_MCTS
from Arena import Arena
class Colors:
RESET = "\033[0m"
RED = "\033[91m"
GREEN = "\033[92m"
YELLOW = "\033[93m"
BLUE = "\033[94m"
MAGENTA = "\033[95m"
CYAN = "\033[96m"
WHITE = "\033[97m"
class Alpha_Zero:
def __init__(self, game, args, model, optimizer):
self.game = game
self.args = args
self.model = model
self.optimizer = optimizer
self.mcts = Alpha_MCTS(game, args, model)
def self_play(self):
single_game_memory = []
player = 1
state = self.game.initialise_state()
while True:
neutral_state = self.game.change_perspective(state, player) if self.args["ADVERSARIAL"] else state
prob = self.mcts.search(neutral_state)
single_game_memory.append((neutral_state, prob, player))
temp_prob = prob ** (1 / self.args["TEMPERATURE"])
temp_prob[temp_prob == 0] = - np.inf
temp_prob = torch.softmax(torch.tensor(temp_prob), axis = 0).cpu().numpy()
move = np.random.choice(self.game.possible_state, p = temp_prob)
state = self.game.make_move(state, move, player)
is_terminal, value = self.game.know_terminal_value(state, move)
if is_terminal:
return_memory = []
for return_state, return_action_prob, return_player in single_game_memory:
if self.args["ADVERSARIAL"]:
return_value = value if return_player == player else self.game.get_opponent_value(value)
else:
return_value = value
return_memory.append((
self.game.get_encoded_state(return_state),
return_action_prob,
return_value
))
return return_memory
if self.args["ADVERSARIAL"]:
player = self.game.get_opponent(player)
def train(self, memory):
random.shuffle(memory)
for batch_start in range(0, len(memory), self.args["BATCH_SIZE"]):
batch_end = batch_start + self.args["BATCH_SIZE"]
training_memory = memory[batch_start : batch_end]
state, action_prob, value = zip(*training_memory)
state, action_prob, value = np.array(state), np.array(action_prob), np.array(value).reshape(-1, 1)
state = torch.tensor(state, device = self.model.device, dtype=torch.float32)
policy_targets = torch.tensor(action_prob, device = self.model.device, dtype=torch.float32)
value_targets = torch.tensor(value, device = self.model.device, dtype=torch.float32)
out_policy, out_value = self.model(state)
policy_loss = F.cross_entropy(out_policy, policy_targets)
value_loss = F.mse_loss(out_value, value_targets)
loss = policy_loss + value_loss
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
def learn(self):
try:
model_path = os.path.join(self.args["MODEL_PATH"], 'model.pt')
optimizer_path = os.path.join(self.args["MODEL_PATH"], 'optimizer.pt')
self.model.load_state_dict(torch.load(model_path))
self.optimizer.load_state_dict(torch.load(optimizer_path))
except:
print(Colors.RED + "UNABLE TO LOAD MODEL")
print(Colors.GREEN + "SETTING UP NEW MODEL..." + Colors.RESET)
else:
print(Colors.GREEN + "MODEL FOUND\nLOADING MODEL..." + Colors.RESET)
finally:
initial_model = copy.copy(self.model)
for iteration in range(self.args["NO_ITERATIONS"]):
memory = []
print(Colors.BLUE + "\nIteration no: " , iteration + 1, Colors.RESET)
print(Colors.YELLOW + "Self Play" + Colors.RESET)
self.model.eval()
for _ in trange(self.args["SELF_PLAY_ITERATIONS"]):
memory += self.self_play()
print(Colors.YELLOW + "Training..." + Colors.RESET)
self.model.train()
for _ in trange(self.args["EPOCHS"]):
self.train(memory)
print(Colors.YELLOW + "Testing..." + Colors.RESET)
self.model.eval()
wins, draws, defeats = Arena(self.game, self.args, self.model, initial_model)
print(Colors.GREEN + "Testing Completed" + Colors.WHITE + "\nTrained Model Stats:")
print(Colors.GREEN, "Wins: ", wins, Colors.RESET, "|", Colors.RED, "Loss: ", defeats, Colors.RESET, "|", Colors.WHITE," Draw: ", draws, Colors.RESET)
print(Colors.YELLOW + "Saving Model...")
torch.save(self.model.state_dict(), os.path.join(self.args["MODEL_PATH"], "model_non_parallel.pt"))
torch.save(self.optimizer.state_dict(), os.path.join(self.args["MODEL_PATH"], "optimizer_non_parallel.pt"))
print("Saved!" + Colors.RESET)
|