Move softmax out of network

[adamantivm]

It also changes the order of masking and applies masking before softmax, like the openspiel implementation.

as a result of this change, the loss now plateaus at around 0.55

[jonbinney]

Looks like the right change. I think I've sorted out the last of my confusion about torch's cross_entropy. I tried testing it this morning by doing:

>>> x = torch.Tensor([0.2, 0.8])
>>> F.cross_entropy(torch.log(x), x)
tensor(0.5004)

This confused me because I thought that the cross entropy of a distribution with itself would be zero, but that's not the case. I used the math expression for torch's cross entropy in the case where the target are probabilities (not labels) from https://docs.pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html and tried replacing the input with the log of the output:

The result is that the cross entropy for a distribution and itself is equal to the entropy of the distribution. Since [0.2, 0.8] has inherent uncertainty, the cross entropy loss will never be zero. That could also explain what you're seeing. If MCTS never produces a 1 and all zeros, the cross entropy loss will never go to zero no matter how well the NN models it.

I'm still a bit confused about why we won't end up with negative infinities in the logits coming out of the network if it tries real hard to match a target with an index that always has zero in some indices because of masking. Maybe torch's cross entropy loss magically fixes this somehow?

EDIT: produces => never produces

[jonbinney]

LGTM. I'll PR my test for training once this is in.

[jonbinney]

Side note: since a perfect NN will have a policy loss of at least the entropy of the target distribution, the (cross_entropy(logit, target) - entropy(target) might be a good metric to track for the training.

[adamantivm]

Very interesting analysis @jonbinney , thank you. I was also under the wrong impression that the loss should be 0 when it converges. It's much better to have this crisper understanding. It will take a while to commit this understanding to intuition, but now there is no arguiing with this.