update pg、predict and add rl

rl/ac.py

@@ -0,0 +1,110 @@
+import numpy as np
+import torch
+import torch.optim as optim
+
+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 ACAgent(Agent):
+    def __init__(self, model, encoder):
+        super().__init__()
+        self.model = model
+        self.encoder = encoder
+        self.collector = None
+        self.temperature = 1.0
+        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        self.model = self.model.to(self.device)
+        self.last_state_value = 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], dtype=torch.float32).to(self.device)
+
+        actions, values = self.model(x)
+        move_probs = actions[0].cpu().detach().numpy()
+        estimated_value = values[0][0].cpu().detach().numpy()
+        self.last_state_value = float(estimated_value)
+
+        # Prevent move probs from getting stuck at 0 or 1.
+        move_probs = np.power(move_probs, 1.0 / self.temperature)
+        move_probs = move_probs / np.sum(move_probs)
+        eps = 1e-6
+        move_probs = np.clip(move_probs, eps, 1 - eps)
+        # Re-normalize to get another probability distribution.
+        move_probs = move_probs / np.sum(move_probs)
+
+        # Turn the probabilities into a ranked list of moves.
+        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,
+                        estimated_value=estimated_value
+                    )
+                return goboard.Move.play(point)
+        # No legal, non-self-destructive moves less.
+        return goboard.Move.pass_turn()
+
+    def train(self, experience, lr=0.1, batch_size=128):
+        opt = optim.Adam(self.model.parameters(), lr=lr)
+
+        n = experience.states.shape[0]
+        num_moves = self.encoder.num_points()
+        policy_target = np.zeros((n, num_moves))
+        value_target = np.zeros((n,))
+        for i in range(n):
+            action = experience.actions[i]
+            reward = experience.rewards[i]
+            policy_target[i][action] = experience.advantages[i]
+            value_target[i] = reward
+
+        policy_target = torch.tensor(policy_target, dtype=torch.float32).to(self.device)
+        value_target = torch.tensor(value_target, dtype=torch.float32).to(self.device)
+        states = torch.tensor(experience.states, dtype=torch.float32).to(self.device)
+
+        for _ in range(n // batch_size):
+            self.model.train()
+            opt.zero_grad()
+            action_pred, value_pred = self.model(states)
+            value_loss = ((value_target - value_pred) ** 2).mean()
+            policy_loss = -(policy_target * torch.log(action_pred)).sum(1).mean()
+            loss = policy_loss + value_loss
+            loss.backward()
+            opt.step()
+
+    def diagnostics(self):
+        return {'value': self.last_state_value}
+
+    def save(self, path):
+        torch.save(self.model.state_dict(), path)
+
+    @classmethod
+    def load(cls, 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(path))
+        return cls(model, encoder)
+
+# def load_ac_agent(path, encoder):
+#     # 在这个转换中，MyModel() 是你的 PyTorch 模型的实例化方法，你需要将其替换为你实际使用的模型。并且，假设你的模型接受一个状态并输出行动和值。
+#     model = MyModel()  # Initialize your model
+#     model.load_state_dict(torch.load(path))
+#     return ACAgent(model, encoder)