Simple policy gradients (REINFORCE) overfits one action when playing Atari Breakout
Self-contained code: https://colab.research.google.com/drive/1HYEXMpicymPUySkhGOaCJdJ3pN4RzXYd

The problem: I am trying to use a CNN to play Atari breakout from pixels (Breakout-V0 OpenAI gym). I am trying to use the simple policy gradients algorithm, implemented in PyTorch, to do this. There are four possible actions in this game [&lt;NO-OP&gt;, &lt;FIRE&gt; (play), &lt;LEFT&gt;, &lt;RIGHT&gt;].

Expected results: I would expect the policy to learn to play &lt;NO-OP&gt;, &lt;LEFT&gt;, &lt;RIGHT&gt; with about equal probability, and only play &lt;FIRE&gt; on the first frame of the game.

Actual results: After ~4 weight updates, the network predicts ONE action with nearly 100% probability. This means the gradient vanishes and the policy never recovers.

What I've tried: 


Only play random actions at the start of the game (take increasingly greedy actions that follow the policy). You can fiddle around with this in the notebook.
I have rewarded only one action, just to prove to myself that it learns to only play that one action. It does, so I think I can rule out any PyTorch specific implementation errors.
Introducing an entropy penalty to the loss to discourage high confidence in actions.


My understanding is that this should not be necessary to explicitly introduce random actions because action = categorical.sample() does this. And if one action becomes dominant, but does not lead to reward, it should be subsequently discouraged.

My thought was that my training batches were skewed so that taking the action &lt;RIGHT&gt; for example lead to more reward per episode than punishment, and hence it's likelihood kept increasing. I would have thought that taking totally random actions at the start of the game and only slowly beginning to listen to the policy would have fixed this, but in my experiments, it didn't.

I'm really very confused as to why this doesn't work. BIG thanks to anyone who can help. I tried to debug this with a PhD in RL for 5 hours yesterday and made no progress.

Extra Questions:


Is it common to play e.g. 1000 rollouts, then sample a batch randomly from this rollout buffer to learn from? My understanding of policy gradients was that the policy is updated after every episode.
Am I correct in thinking that in policy gradients, you only backprop through the neuron that you selected your action from, but the gradient gets distributed to ALL the network weights through the softmax?


We ran into a similar problem while training Breakout with VPG (Vanilla Policy Gradient). The solution was to enforce entropy loss over the following


Entropy loss over the policy model output (Penalise selecting a specific action with high likelihood)
Scaling the MSE between rewards to go and value function


Not sure if you are aware of this but do use the Atari wrapper from Deepmind

Q 1 : Is it common to play e.g. 1000 rollouts, then sample a batch randomly from this rollout buffer to learn from? My understanding of policy gradients was that the policy is updated after every episode.

A 1 : In practice its better to batch multiple episodes in a single batch.

Q 2 : Am I correct in thinking that in policy gradients, you only backprop through the neuron that you selected your action from, but the gradient gets distributed to ALL the network weights through the softmax?

A 2 : Based on the expected rewards to go, likelihood of selection of specific action is suppressed or encouraged and the model backdrops this information
