I'm having some difficulty applying policy gradient to the simple OpenAI Gym environment CartPole-v0. This is a simple inverted pendulum where the agent receives a reward of 1 in every frame where the pendulum is within a threshold angle and position.
Learning seems to start well, with the average reward per episode steadily increasing, until a catastrophic event occurs and the reward collapses and never recovers (see image below, 8 trials shown). Is this a common thing to encounter when using policy gradient? Any idea what is going wrong and how to fix it? Code here
More details: I predict the action logits as a linear function of the state, and I use a one-hidden-layer network to estimate the state-value function to obtain the advantage (i.e. A2C), similar to this report. Discount factor gamma is 0.9. I use 30 episodes to compute the gradient each time. I'm using SGD with learning rate 1e-8 (low?) and momentum 0.9. Implementation is Tensorflow. I have a maximum episode length of 1000 time steps. I have added a quadratic regularizer with weight 1.
Magnitude of gradients
Here I visualize the max-norm of each parameter on the second (right) y-axis with dotted lines. The gradients seem enormous, so maybe this is the problem, although they seem quite large before the collapse too. Here are two trials:
Params (0, 1) are (weight, bias) for the single linear layer that predicts the action logits from the state. Params (2, 3) and (4, 5) are (weight, bias) for the two-layer network that estimates V(s).