Need help about my Connect 4 Deep Reinforcement Learning Agent - machine-learning

Question

Here is my code
My env setting
state: np.array, 1 * 6 * 7 (1 input channel for CNN, 6 * 7 the size of the board), 0 = nothing, 1 = agent1's token, -1 =
agent2's token
reward function: 0 if draw or nothing happen, 1 if win, -1 if lose
I know that Connect 4 is a solved game, but I want to try making two value based method agents to learn the game from scratch. I tried ANN with flatten the state as input, and output 7 value(argmax to get action), but this method make the agent converge to a stupid strategy.
Example
_ | _ | _ | _ | _ | _ | _
_ | _ | _ | _ | _ | _ | _
_ | X | _ | _ | _ | _ | _
_ | X | _ | O | _ | _ | _
_ | X | _ | O | _ | _ | _
_ | X | _ | O | _ | _ | _
Then the agent who go first always win with this "strategy".
Then, I tried CNN and n step learning but it still converge to the "strategy", here is the CNN architecture
class Network(nn.Module):
def __init__(self, output_dim: int, learning_rate: float) -> None:
super(Network, self).__init__()
self.feature_layers = nn.Sequential(
nn.Conv2d(1, 64, 4),
nn.ReLU(),
nn.Conv2d(64, 32, 2),
nn.ReLU(),
nn.Conv2d(32, 16, 2),
nn.Flatten(),
)
self.advantage_layers = nn.Sequential(
nn.Linear(32, 16),
nn.ReLU(),
nn.Linear(16, output_dim)
)
self.value_layers = nn.Sequential(
nn.Linear(32, 8),
nn.ReLU(),
nn.Linear(8, 1)
)
self.optimizer = optim.RMSprop(self.parameters(), lr = learning_rate)
self.loss = nn.MSELoss()
def forward(self, x: torch.Tensor) -> torch.Tensor:
feature = self.feature_layers(x)
advantage = self.advantage_layers(feature)
value = self.value_layers(feature)
return value + advantage - advantage.mean(dim = -1, keepdim = True)
Is there anything I did wrong
(I know that MCTS is very good for this case, but I am still learning value based method)