Why training baseline of REINFORCE by MSE?

Question

In following references

 [1] V. Mnih et al., "Recurrent Models of Visual Attention", NIPS, 2014
 [2] http://torch.ch/blog/2015/09/21/rmva.html
 [3] M. Ranzato et al., "Sequence level training with Recurrent Neural Network", ICLR, 2016

they applied REINFORCE algorithm to train RNN.

To reduce variance of the gradient, they subtract 'baseline' from sum of future rewards for all time steps. According to Appendix A-2 of

[4]. W. Zaremba et al., "Reinforcement Learning Neural Turing Machines", arXiv, 2016

this baseline is chosen as expected future reward given previous states/actions.

My question is training method to get 'baseline'. They train 'baseline' at each time by linear regression (i.e. objective = Mean Square Error) which takes hidden state of RNN as input. How come this training method make sense? Is there anyone to provide relevant reference?

I really appreciate for your comments.

Answer 1

So the reason you want a baseline is that $R_\tau$ can be very different from one episode to another under the policy you are training (equation no 3 in the paper). Subtracting a baseline from that return helps reducing the variance and stabilizing the algorithm. See for example John Schulman's MLSS Lecture or David Silver's lecture on policy gradient methods.

When training policy gradient methods the baseline can be whatever approximates best the expected reward onward from the step you are in. A natural baseline is the state value function, the value of being in that state $V(s)$. Usually these functions are learned by parametric estimators (linear function approximators, neural networks, etc.). More on the general theory of using baselines in policy gradient methods you can find here.

In the paper you linked they are using a separate LSTM that learns to predict the expected episodic return from that step onward, $R_\tau$ so it is only natural that the objective to be minimized is the squared mean error between the actual return $R_\tau$ and the value predicted by the baseline LSTM, $b_\tau$.

The features the baseline LSTM is using are the full tape at the first time step and then the exact input the RL-LSTM is receiving at each step which is a fairly complex statistic (past actions, current memory, current input, hidden state of RL-LSTM). If you are asking if this is a sufficient statistic to predict the episodic return $R_\tau$, well, it seems to be working :). For example a similar baseline was learned in this attention model + implementation.