Hi,
three things I noticed (but I’m not an expert in RL)

I would expect update_parameters to be called less often. Should this be moved outside the for loop? I seem to associate policy gradient with notoriously high variance, so deriving one update from multiple trials might work better.
I would clamp the probabilities away from zero before taking logs. I needed this to avoid the device asserts.
Are you sure the “R” formula does what you want here? For a bandit type prediction, this seems somewhat odd.

Changing these, the network appears to learn something (being mostly right or off only by 1) although the results won’t be perfect even after many iterations.
Best regards
Thomas