some try for policy_gradient

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)