fix(policies): keep invalid actions masked during sampling

[Linxiushen]

At high temperatures, muzero_policy and stochastic_muzero_policy can return an action marked invalid even though the search never visits it. Zero visit probabilities are clipped to the dtype's finite tiny value before taking logs; temperature scaling then gives these actions nonzero sampling probability. With temperature 100 and one valid action per row, a deterministic 32-row example returns 17 invalid actions.

Reapply the existing invalid-action mask after temperature scaling and before the final categorical sample in both policies. Search results and the returned training weights remain unchanged.

The regressions run real MuZero and stochastic MuZero searches, with a different valid action across rows, both directly and under JIT, at temperatures 0, 1, and 100. On unchanged main, the four high-temperature cases fail and eight controls pass. With the fix, the full test suite passes: 36 passed on Windows Python 3.12.13 / JAX 0.11.2 CPU, run from the _testing directory prescribed by test.sh.

The CI flake8 selector passes repository-wide, production-file Pylint passes, and wheel/sdist builds succeed. The configured test Pylint and Ruff checks retain one pre-existing long-line diagnostic each, outside the modified code; no new findings. Pyink passes for the added test block. No GPU/TPU, training run, Linux run, or supported-Python matrix is claimed.

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.