From 41c70802fd805fbebc7801449155b296a392ef87 Mon Sep 17 00:00:00 2001
From: bbig <huangtengle@gmail.com>
Date: Tue, 30 May 2023 18:48:27 +0800
Subject: [PATCH] =?UTF-8?q?update=20pg=E3=80=81predict=20and=20add=20rl?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 agents/policy_gradient_agent.py |   3 +-
 agents/predict.py               |   3 +-
 rl/__init__.py                  |   0
 rl/ac.py                        | 110 +++++++++++++++++++++++
 rl/ac_pass.py                   |   0
 rl/experience.py                |  70 +++++++++++++++
 rl/q.py                         | 155 ++++++++++++++++++++++++++++++++
 rl/simulate.py                  |   0
 rl/value.py                     | 149 ++++++++++++++++++++++++++++++
 9 files changed, 488 insertions(+), 2 deletions(-)
 create mode 100644 rl/__init__.py
 create mode 100644 rl/ac.py
 create mode 100644 rl/ac_pass.py
 create mode 100644 rl/experience.py
 create mode 100644 rl/q.py
 create mode 100644 rl/simulate.py
 create mode 100644 rl/value.py

diff --git a/agents/policy_gradient_agent.py b/agents/policy_gradient_agent.py
index 3dcefd8..c9aa74f 100644
--- a/agents/policy_gradient_agent.py
+++ b/agents/policy_gradient_agent.py
@@ -86,6 +86,7 @@ class PolicyAgent(Agent):
 
     @classmethod
     def load(cls, path, encoder):
-        model = AlphaGoModel(encoder.get_input_shape(), is_policy_net=True).to('cuda')
+        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)
diff --git a/agents/predict.py b/agents/predict.py
index 3b3bd27..77f2a13 100644
--- a/agents/predict.py
+++ b/agents/predict.py
@@ -43,6 +43,7 @@ class DeepLearningAgent(Agent):
 
     @classmethod
     def load(cls, path, encoder):
-        model = AlphaGoModel(encoder.get_input_shape(), is_policy_net=True).to('cuda')
+        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)
diff --git a/rl/__init__.py b/rl/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/rl/ac.py b/rl/ac.py
new file mode 100644
index 0000000..61da5af
--- /dev/null
+++ b/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)
diff --git a/rl/ac_pass.py b/rl/ac_pass.py
new file mode 100644
index 0000000..e69de29
diff --git a/rl/experience.py b/rl/experience.py
new file mode 100644
index 0000000..4946e7f
--- /dev/null
+++ b/rl/experience.py
@@ -0,0 +1,70 @@
+import torch
+import pickle
+
+class ExperienceCollector:
+    def __init__(self):
+        self.states = []
+        self.actions = []
+        self.rewards = []
+        self.advantages = []
+        self._current_episode_states = []
+        self._current_episode_actions = []
+        self._current_episode_estimated_values = []
+
+    def begin_episode(self):
+        self._current_episode_states = []
+        self._current_episode_actions = []
+        self._current_episode_estimated_values = []
+
+    def record_decision(self, state, action, estimated_value=0):
+        self._current_episode_states.append(state)
+        self._current_episode_actions.append(action)
+        self._current_episode_estimated_values.append(estimated_value)
+
+    def complete_episode(self, reward):
+        num_states = len(self._current_episode_states)
+        self.states += self._current_episode_states
+        self.actions += self._current_episode_actions
+        self.rewards += [reward for _ in range(num_states)]
+
+        for i in range(num_states):
+            advantage = reward - self._current_episode_estimated_values[i]
+            self.advantages.append(advantage)
+
+        self._current_episode_states = []
+        self._current_episode_actions = []
+        self._current_episode_estimated_values = []
+
+
+class ExperienceBuffer:
+    def __init__(self, states, actions, rewards, advantages):
+        self.states = states
+        self.actions = actions
+        self.rewards = rewards
+        self.advantages = advantages
+
+    def save(self, filename):
+        with open(filename, 'wb') as f:
+            pickle.dump((self.states, self.actions, self.rewards, self.advantages), f)
+
+    @classmethod
+    def load(cls, filename):
+        with open(filename, 'rb') as f:
+            states, actions, rewards, advantages = pickle.load(f)
+        return cls(
+            states=states,
+            actions=actions,
+            rewards=rewards,
+            advantages=advantages)
+
+def combine_experience(collectors):
+    combined_states = torch.cat([torch.tensor(c.states) for c in collectors])
+    combined_actions = torch.cat([torch.tensor(c.actions) for c in collectors])
+    combined_rewards = torch.cat([torch.tensor(c.rewards) for c in collectors])
+    combined_advantages = torch.cat([torch.tensor(c.advantages) for c in collectors])
+
+    return ExperienceBuffer(
+        combined_states,
+        combined_actions,
+        combined_rewards,
+        combined_advantages)
diff --git a/rl/q.py b/rl/q.py
new file mode 100644
index 0000000..9f34a15
--- /dev/null
+++ b/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)
diff --git a/rl/simulate.py b/rl/simulate.py
new file mode 100644
index 0000000..e69de29
diff --git a/rl/value.py b/rl/value.py
new file mode 100644
index 0000000..2af6e44
--- /dev/null
+++ b/rl/value.py
@@ -0,0 +1,149 @@
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.distributions import Categorical
+
+from dlgo import encoders
+from dlgo import goboard
+from dlgo.agent import Agent
+from dlgo.agent.helpers import is_point_an_eye
+
+
+class ValueAgent(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 predict(self, game_state):
+        encoded_state = self.encoder.encode(game_state)
+        input_tensor = torch.tensor([encoded_state])
+        return self.model(input_tensor)[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):
+        moves = []
+        board_tensors = []
+        for move in game_state.legal_moves():
+            if not move.is_play:
+                continue
+            next_state = game_state.apply_move(move)
+            board_tensor = self.encoder.encode(next_state)
+            moves.append(move)
+            board_tensors.append(board_tensor)
+        if not moves:
+            return goboard.Move.pass_turn()
+
+        board_tensors = torch.tensor(board_tensors)
+
+        opp_values = self.model(board_tensors)
+        opp_values = opp_values.reshape(len(moves))
+
+        values = 1 - opp_values
+
+        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:
+            move = moves[move_idx]
+            if not is_point_an_eye(game_state.board,
+                                   move.point,
+                                   game_state.next_player):
+                if self.collector is not None:
+                    self.collector.record_decision(
+                        state=board_tensor,
+                        action=self.encoder.encode_point(move.point),
+                    )
+                self.last_move_value = float(values[move_idx])
+                return move
+        return goboard.Move.pass_turn()
+
+    def rank_moves_eps_greedy(self, values):
+        if torch.rand(1).item() < self.temperature:
+            values = torch.rand_like(values)
+        ranked_moves = torch.argsort(values)
+        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]
+        y = torch.zeros((n,))
+        for i in range(n):
+            reward = experience.rewards[i]
+            y[i] = 1 if reward > 0 else 0
+
+        self.model.train()
+        opt.zero_grad()
+        pred = self.model(experience.states)
+        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)
+
+        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_value_agent(filename):
+    # 我假设 Model 类是预先定义好的 PyTorch 模型，你需要根据实际情况替换或实现这个模型。
+    checkpoint = torch.load(filename)
+
+    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 ValueAgent(model, encoder)