From e0ef39bc2766e6b1de5be73c26d962fb4eb35ebc Mon Sep 17 00:00:00 2001 From: Julian Cerruti Date: Fri, 11 Jul 2025 17:19:51 -0300 Subject: [PATCH 1/3] Move softmax out of network --- deep_quoridor/src/agents/alphazero/mlp_network.py | 5 ++++- deep_quoridor/src/agents/alphazero/nn_evaluator.py | 7 ++++--- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/deep_quoridor/src/agents/alphazero/mlp_network.py b/deep_quoridor/src/agents/alphazero/mlp_network.py index 3fc83cb5..d14fa59b 100644 --- a/deep_quoridor/src/agents/alphazero/mlp_network.py +++ b/deep_quoridor/src/agents/alphazero/mlp_network.py @@ -24,7 +24,10 @@ def __init__(self, input_size, action_size, device): # TODO: Is it correct to include the Softmax at the end? Some implementations of alphazero # appear to leave it out, or apply it outside the network. self.policy_head = nn.Sequential( - nn.Linear(256, 128), nn.ReLU(), nn.Linear(128, action_size), nn.Softmax(dim=-1) + # nn.Linear(256, 128), nn.ReLU(), nn.Linear(128, action_size), nn.Softmax(dim=-1) + nn.Linear(256, 128), + nn.ReLU(), + nn.Linear(128, action_size), ) # value head - outputs position evaluation (-1 to 1) diff --git a/deep_quoridor/src/agents/alphazero/nn_evaluator.py b/deep_quoridor/src/agents/alphazero/nn_evaluator.py index 31ef558e..bb300f61 100644 --- a/deep_quoridor/src/agents/alphazero/nn_evaluator.py +++ b/deep_quoridor/src/agents/alphazero/nn_evaluator.py @@ -38,7 +38,8 @@ def evaluate(self, game: Quoridor): with torch.no_grad(): input_array = self.game_to_input_array(game) - unmasked_policy, value = self.network(torch.from_numpy(input_array).float().to(self.device)) + unmasked_policy_logits, value = self.network(torch.from_numpy(input_array).float().to(self.device)) + unmasked_policy = F.softmax(unmasked_policy_logits, dim=-1) unmasked_policy = unmasked_policy.cpu().numpy() value = value.item() @@ -159,10 +160,10 @@ def train_iteration(self, replay_buffer): raise ValueError("NaN in training data") # Forward pass - pred_policies, pred_values = self.network(inputs) + pred_logits, pred_values = self.network(inputs) # Compute losses - policy_loss = F.cross_entropy(pred_policies, target_policies, reduction="mean") + policy_loss = F.cross_entropy(pred_logits, target_policies, reduction="mean") value_loss = F.mse_loss(pred_values.squeeze(), target_values.squeeze(), reduction="mean") total_loss = policy_loss + value_loss # print(f"{total_loss.item():3.3f} {policy_loss.item():3.3f} {value_loss.item():3.3f}") From 30f7de2197908b5f1517e3a1d4a553a72aaf6288 Mon Sep 17 00:00:00 2001 From: Julian Cerruti Date: Fri, 11 Jul 2025 17:44:20 -0300 Subject: [PATCH 2/3] Apply masking before softmax --- .../src/agents/alphazero/nn_evaluator.py | 29 ++++++++----------- 1 file changed, 12 insertions(+), 17 deletions(-) diff --git a/deep_quoridor/src/agents/alphazero/nn_evaluator.py b/deep_quoridor/src/agents/alphazero/nn_evaluator.py index bb300f61..74ae178e 100644 --- a/deep_quoridor/src/agents/alphazero/nn_evaluator.py +++ b/deep_quoridor/src/agents/alphazero/nn_evaluator.py @@ -36,27 +36,21 @@ def evaluate(self, game: Quoridor): self.network.eval() # Disables dropout + valid_actions = game.get_valid_actions() + valid_action_indices = [self.action_encoder.action_to_index(action) for action in valid_actions] + with torch.no_grad(): input_array = self.game_to_input_array(game) unmasked_policy_logits, value = self.network(torch.from_numpy(input_array).float().to(self.device)) - unmasked_policy = F.softmax(unmasked_policy_logits, dim=-1) - unmasked_policy = unmasked_policy.cpu().numpy() - value = value.item() + unmasked_policy_logits = unmasked_policy_logits.cpu().numpy() - # Mask the policy to ignore invalid actions. NOTE: Game is already rotated so the valid actions will be rotated too - valid_actions = game.get_valid_actions() - valid_action_indices = [self.action_encoder.action_to_index(action) for action in valid_actions] - policy_masked = np.zeros_like(unmasked_policy) - policy_masked[valid_action_indices] = unmasked_policy[valid_action_indices] - - if np.all(policy_masked == 0): - # If the policy ends up as all zeros after masking, turn it into a uniform distribution among - # the valid actions. - policy_masked[valid_action_indices] = 1 / len(valid_action_indices) - print("Policy is all zeros after masking, turning it into a uniform distribution") - else: - # Otherwise, just renormalize after masking - policy_masked = policy_masked / policy_masked.sum() + # Set logits of invalid actions to negative infinity to ensure they are never chosen + # TODO: Mask entirely in Torch to avoid going to Numpy and back + masked_logits = np.full_like(unmasked_policy_logits, -np.inf) + masked_logits[valid_action_indices] = unmasked_policy_logits[valid_action_indices] + + policy_masked = F.softmax(torch.from_numpy(masked_logits), dim=-1).cpu().numpy() + value = value.item() # Sanity checks assert np.all(policy_masked >= 0), "Policy contains negative probabilities" @@ -161,6 +155,7 @@ def train_iteration(self, replay_buffer): # Forward pass pred_logits, pred_values = self.network(inputs) + # TODO: Should we apply masking before calculating cross-entropy here? # Compute losses policy_loss = F.cross_entropy(pred_logits, target_policies, reduction="mean") From 81c9efd9f10bebc592b00a5b2dae57ab956bb115 Mon Sep 17 00:00:00 2001 From: Julian Cerruti Date: Fri, 11 Jul 2025 17:44:57 -0300 Subject: [PATCH 3/3] remove unnecessary comment --- deep_quoridor/src/agents/alphazero/mlp_network.py | 1 - 1 file changed, 1 deletion(-) diff --git a/deep_quoridor/src/agents/alphazero/mlp_network.py b/deep_quoridor/src/agents/alphazero/mlp_network.py index d14fa59b..f2f94a00 100644 --- a/deep_quoridor/src/agents/alphazero/mlp_network.py +++ b/deep_quoridor/src/agents/alphazero/mlp_network.py @@ -24,7 +24,6 @@ def __init__(self, input_size, action_size, device): # TODO: Is it correct to include the Softmax at the end? Some implementations of alphazero # appear to leave it out, or apply it outside the network. self.policy_head = nn.Sequential( - # nn.Linear(256, 128), nn.ReLU(), nn.Linear(128, action_size), nn.Softmax(dim=-1) nn.Linear(256, 128), nn.ReLU(), nn.Linear(128, action_size),