1

agent/predict.py

@@ -0,0 +1,57 @@
+# tag::dl_agent_imports[]
+import numpy as np
+
+from tugo.agent.base import Agent
+from tugo.agent.helpers import is_point_an_eye
+from tugo import encoders
+from tugo import goboard
+from tugo import kerasutil
+from tugo.networks import AlphaGoModel
+# end::dl_agent_imports[]
+__all__ = [
+    'DeepLearningAgent',
+]
+
+import numpy as np
+import torch
+
+class DeepLearningAgent(Agent):
+    def __init__(self, model, encoder):
+        super().__init__()
+        self.model = model
+        self.encoder = encoder
+
+    def predict(self, game_state):
+        encoded_state = self.encoder.encode(game_state)
+        input_tensor = torch.tensor([encoded_state], dtype=torch.float32).to('cuda')
+        with torch.no_grad():
+            output_tensor = self.model(input_tensor)
+        return output_tensor.cpu().numpy()[0]
+
+    def select_move(self, game_state):
+        num_moves = self.encoder.board_width * self.encoder.board_height
+        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):
+                return goboard.Move.play(point)
+        return goboard.Move.pass_turn()
+
+    def save(self, path):
+        torch.save(self.model.state_dict(), path)
+
+    @classmethod
+    def load(cls, path, encoder):
+        model = AlphaGoModel(encoder.get_input_shape(), is_policy_net=True).to('cuda')
+        model.load_state_dict(torch.load(path))
+        return cls(model, encoder)