update pg、predict and add rl

rl/q.py

@@ -0,0 +1,155 @@
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.distributions import Categorical
+
+from tugo import encoders
+from tugo.game_logic import goboard
+from tugo.agents import Agent
+from tugo.agents.helpers import is_point_an_eye
+
+class QAgent(Agent):
+    def __init__(self, model, encoder, policy='eps-greedy'):
+        Agent.__init__(self)
+        self.model = model
+        self.encoder = encoder
+        self.collector = None
+        self.temperature = 0.0
+        self.policy = policy
+
+        self.last_move_value = 0
+
+    def set_temperature(self, temperature):
+        self.temperature = temperature
+
+    def set_collector(self, collector):
+        self.collector = collector
+
+    def set_policy(self, policy):
+        if policy not in ('eps-greedy', 'weighted'):
+            raise ValueError(policy)
+        self.policy = policy
+
+    def select_move(self, game_state):
+        board_tensor = self.encoder.encode(game_state)
+
+        # Loop over all legal moves.
+        moves = []
+        board_tensors = []
+        for move in game_state.legal_moves():
+            if not move.is_play:
+                continue
+            moves.append(self.encoder.encode_point(move.point))
+            board_tensors.append(board_tensor)
+        if not moves:
+            return goboard.Move.pass_turn()
+
+        num_moves = len(moves)
+        board_tensors = torch.tensor(board_tensors)
+        move_vectors = torch.zeros((num_moves, self.encoder.num_points()))
+        for i, move in enumerate(moves):
+            move_vectors[i][move] = 1
+
+        values = self.model([board_tensors, move_vectors])
+        values = values.reshape(len(moves))
+
+        if self.policy == 'eps-greedy':
+            ranked_moves = self.rank_moves_eps_greedy(values)
+        elif self.policy == 'weighted':
+            ranked_moves = self.rank_moves_weighted(values)
+        else:
+            ranked_moves = None
+
+        for move_idx in ranked_moves:
+            point = self.encoder.decode_point_index(moves[move_idx])
+            if 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=moves[move_idx],
+                    )
+                self.last_move_value = float(values[move_idx])
+                return goboard.Move.play(point)
+        # No legal, non-self-destructive moves less.
+        return goboard.Move.pass_turn()
+
+    def rank_moves_eps_greedy(self, values):
+        if torch.rand(1).item() < self.temperature:
+            values = torch.rand_like(values)
+        # This ranks the moves from worst to best.
+        ranked_moves = torch.argsort(values)
+        # Return them in best-to-worst order.
+        return ranked_moves[::-1]
+
+    def rank_moves_weighted(self, values):
+        p = values / torch.sum(values)
+        p = torch.pow(p, 1.0 / self.temperature)
+        p = p / torch.sum(p)
+        return Categorical(p).sample().item()
+
+    def train(self, experience, lr=0.1, batch_size=128):
+        opt = optim.SGD(self.model.parameters(), lr=lr)
+        criterion = nn.MSELoss()
+
+        n = experience.states.shape[0]
+        num_moves = self.encoder.num_points()
+        y = torch.zeros((n,))
+        actions = torch.zeros((n, num_moves))
+        for i in range(n):
+            action = experience.actions[i]
+            reward = experience.rewards[i]
+            actions[i][action] = 1
+            y[i] = 1 if reward > 0 else 0
+
+        self.model.train()
+        opt.zero_grad()
+        pred = self.model([experience.states, actions])
+        loss = criterion(pred, y)
+        loss.backward()
+        opt.step()
+
+    def diagnostics(self):
+        return {'value': self.last_move_value}
+
+    def save(self, filename):
+        torch.save({
+            'model': self.model.state_dict(),
+            'encoder': {
+                'name': self.encoder.name(),
+                'board_width': self.encoder.board_width,
+                'board_height': self.encoder.board_height,
+            }
+        }, filename)
+
+    @classmethod
+    def load(cls, filename):
+        checkpoint = torch.load(filename)
+
+        # Assuming `model` is a PyTorch model here and `encoder` is a DLGo encoder
+        model = Model()  # The Model class should be defined appropriately
+        model.load_state_dict(checkpoint['model'])
+
+        encoder_info = checkpoint['encoder']
+        encoder = encoders.get_encoder_by_name(
+            encoder_info['name'],
+            (encoder_info['board_width'], encoder_info['board_height']))
+
+        return cls(model, encoder)
+
+
+def load_q_agent(filename):
+    # 我假设 Model 类是预先定义好的 PyTorch 模型，你需要根据实际情况替换或实现这个模型。
+    checkpoint = torch.load(filename)
+
+    # Assuming `model` is a PyTorch model here and `encoder` is a DLGo encoder
+    model = Model()  # The Model class should be defined appropriately
+    model.load_state_dict(checkpoint['model'])
+
+    encoder_info = checkpoint['encoder']
+    encoder = encoders.get_encoder_by_name(
+        encoder_info['name'],
+        (encoder_info['board_width'], encoder_info['board_height']))
+
+    return QAgent(model, encoder)